Add endpoint for jurisdiction admins to create audit boards

Adds an endpoint
`POST
/election/<election_id>/jurisdiction/<jurisdiction_id>/round/<round_id>/audit-board`
that jurisdiction admins can use to create audit boards.

When the audit boards are created, the sampled ballots for that round
are divvyed up as fairly as possible between the audit boards.

Also adds a db constraint that audit board names must be unique within a
jurisdiction for each round.

arlo_server/__init__.py

@@ -41,12 +41,19 @@
 
 # The order of these imports is important as it defines route precedence.
 # Be careful when re-ordering them.
-import arlo_server.election_settings
+
+# Single-jurisdiction flow routes
+# (Plus some routes for multi-jurisdiction flow that were created before we had
+# separate routes modules)
 import arlo_server.routes
+
+# Multi-jurisdiction flow routes
+import arlo_server.election_settings
 import arlo_server.contests
 import arlo_server.jurisdictions
 import arlo_server.sample_sizes
 import arlo_server.rounds
+import arlo_server.audit_boards
 
 # Error handlers
 import arlo_server.errors

arlo_server/audit_boards.py

@@ -0,0 +1,121 @@
+from flask import jsonify, request
+import uuid
+from typing import List
+from xkcdpass import xkcd_password as xp
+from werkzeug.exceptions import Conflict
+from sqlalchemy.exc import IntegrityError
+
+from arlo_server import app, db
+from arlo_server.auth import with_jurisdiction_access
+from arlo_server.rounds import get_current_round
+from arlo_server.models import (
+    AuditBoard,
+    Round,
+    Election,
+    Jurisdiction,
+    SampledBallot,
+    Batch,
+)
+from arlo_server.errors import handle_unique_constraint_error
+from util.jsonschema import validate, JSONDict
+from util.binpacking import BalancedBucketList, Bucket
+from util.group_by import group_by
+
+WORDS = xp.generate_wordlist(wordfile=xp.locate_wordfile())
+
+CREATE_AUDIT_BOARD_REQUEST_SCHEMA = {
+    "type": "object",
+    "properties": {"name": {"type": "string"},},
+    "additionalProperties": False,
+    "required": ["name"],
+}
+
+# Raises if invalid
+def validate_audit_boards(
+    audit_boards: List[JSONDict],
+    election: Election,
+    jurisdiction: Jurisdiction,
+    round: Round,
+):
+    current_round = get_current_round(election)
+    if not current_round or round.id != current_round.id:
+        raise Conflict(f"Round {round.round_num} is not the current round")
+
+    if any(ab for ab in jurisdiction.audit_boards if ab.round_id == round.id):
+        raise Conflict(f"Audit boards already created for round {round.round_num}")
+
+    validate(
+        audit_boards, {"type": "array", "items": CREATE_AUDIT_BOARD_REQUEST_SCHEMA}
+    )
+
+
+def assign_sampled_ballots(
+    jurisdiction: Jurisdiction, round: Round, audit_boards: List[AuditBoard],
+):
+    # Collect the physical ballots for each batch that were sampled for this
+    # jurisdiction for this round
+    sampled_ballots = (
+        SampledBallot.query.join(Batch)
+        .filter_by(jurisdiction_id=jurisdiction.id)
+        .join(SampledBallot.draws)
+        .filter_by(round_id=round.id)
+        .order_by(SampledBallot.batch_id)  # group_by prefers a sorted list
+        .all()
+    )
+    ballots_by_batch = group_by(sampled_ballots, key=lambda sb: sb.batch_id)
+
+    # Divvy up batches of ballots between the audit boards.
+    # Note: BalancedBucketList doesn't care which buckets have which batches to
+    # start, so we add all the batches to the first bucket before balancing.
+    buckets = [Bucket(audit_board.id) for audit_board in audit_boards]
+    for batch_id, sampled_ballots in ballots_by_batch.items():
+        buckets[0].add_batch(batch_id, len(sampled_ballots))
+    balanced_buckets = BalancedBucketList(buckets)
+
+    for bucket in balanced_buckets.buckets:
+        ballots_in_bucket = [
+            ballot
+            for batch_id in bucket.batches
+            for ballot in ballots_by_batch[batch_id]
+        ]
+        for ballot in ballots_in_bucket:
+            ballot.audit_board_id = bucket.name
+            db.session.add(ballot)
+
+    db.session.commit()
+
+
+@app.route(
+    "/election/<election_id>/jurisdiction/<jurisdiction_id>/round/<round_id>/audit-board",
+    methods=["POST"],
+)
+@with_jurisdiction_access
+def create_audit_boards(election: Election, jurisdiction: Jurisdiction, round_id: str):
+    json_audit_boards = request.get_json()
+    round = Round.query.get_or_404(round_id)
+    validate_audit_boards(json_audit_boards, election, jurisdiction, round)
+
+    audit_boards = [
+        AuditBoard(
+            id=str(uuid.uuid4()),
+            name=json_audit_board["name"],
+            jurisdiction_id=jurisdiction.id,
+            round_id=round.id,
+            passphrase=xp.generate_xkcdpassword(WORDS, numwords=4, delimiter="-"),
+        )
+        for json_audit_board in json_audit_boards
+    ]
+    db.session.add_all(audit_boards)
+
+    try:
+        db.session.commit()
+    except IntegrityError as e:
+        handle_unique_constraint_error(
+            e,
+            constraint_name="audit_board_jurisdiction_id_round_id_name_key",
+            message="Audit board names must be unique",
+        )
+
+    assign_sampled_ballots(jurisdiction, round, audit_boards)
+
+    return jsonify(status="ok")

arlo_server/models.py

@@ -253,6 +253,8 @@ class AuditBoard(BaseModel):
         "SampledBallot", backref="audit_board", passive_deletes=True
     )
 
+    __table_args__ = (db.UniqueConstraint("jurisdiction_id", "round_id", "name"),)
+
 
 class Round(BaseModel):
     id = db.Column(db.String(200), primary_key=True)

arlo_server/routes.py

@@ -1,4 +1,4 @@
-import os, datetime, csv, io, json, uuid, re, hmac, urllib.parse, itertools
+import os, datetime, csv, io, json, uuid, re, hmac, urllib.parse
 
 from flask import jsonify, request, Response, redirect, session
 from flask_httpauth import HTTPBasicAuth
@@ -40,6 +40,7 @@
 from util.isoformat import isoformat
 from util.jsonschema import validate
 from util.process_file import serialize_file, serialize_file_processing
+from util.group_by import group_by
 
 AUDIT_BOARD_MEMBER_COUNT = 2
 WORDS = xp.generate_wordlist(wordfile=xp.locate_wordfile())
@@ -1074,16 +1075,14 @@ def audit_report(election_id):
         # First group all the ballot draws by the actual ballot
         for _, ballot_draws in group_by(
             all_sampled_ballot_draws, key=lambda b: (b.batch_id, b.ballot_position)
-        ):
-            ballot_draws = list(ballot_draws)
+        ).items():
             b = ballot_draws[0]
 
             # Then group the draws for this ballot by round
             ticket_numbers = []
             for round_num, round_draws in group_by(
                 ballot_draws, key=lambda b: b.round.round_num
-            ):
-                round_draws = list(round_draws)
+            ).items():
                 ticket_numbers_str = ", ".join(
                     sorted(d.ticket_number for d in round_draws)
                 )
@@ -1108,10 +1107,6 @@ def audit_report(election_id):
     return response
 
 
-def group_by(xs, key=None):
-    return itertools.groupby(sorted(xs, key=key), key=key)
-
-
 def pretty_affiliation(affiliation):
     mapping = {
         "DEM": "Democrat",

tests/conftest.py

@@ -45,23 +45,29 @@ def election_id(client: FlaskClient) -> Generator[str, None, None]:
 def jurisdiction_ids(
     client: FlaskClient, election_id: str
 ) -> Generator[List[str], None, None]:
+    # We expect the API to order the jurisdictions by name, so we upload them
+    # out of order.
     rv = client.put(
         f"/election/{election_id}/jurisdiction/file",
         data={
             "jurisdictions": (
                 io.BytesIO(
                     b"Jurisdiction,Admin Email\n"
-                    b"J1,a1@example.com\n"
                     b"J2,a2@example.com\n"
-                    b"J3,a3@example.com"
+                    b"J3,a3@example.com\n"
+                    b"J1,a1@example.com"
                 ),
                 "jurisdictions.csv",
             )
         },
     )
     assert_ok(rv)
     bgcompute_update_election_jurisdictions_file()
-    jurisdictions = Jurisdiction.query.filter_by(election_id=election_id).all()
+    jurisdictions = (
+        Jurisdiction.query.filter_by(election_id=election_id)
+        .order_by(Jurisdiction.name)
+        .all()
+    )
     yield [j.id for j in jurisdictions]
 
 

tests/helpers.py

@@ -71,6 +71,17 @@ def create_org_and_admin(
     return org.id, aa.id
 
 
+def create_jurisdiction_admin(
+    jurisdiction_id: str, user_email: str = DEFAULT_JA_EMAIL
+) -> str:
+    ja = create_user(user_email)
+    db.session.add(ja)
+    admin = JurisdictionAdministration(user_id=ja.id, jurisdiction_id=jurisdiction_id)
+    db.session.add(admin)
+    db.session.commit()
+    return str(ja.id)
+
+
 def create_jurisdiction_and_admin(
     election_id: str,
     jurisdiction_name: str = "Test Jurisdiction",
@@ -79,13 +90,10 @@ def create_jurisdiction_and_admin(
     jurisdiction = Jurisdiction(
         id=str(uuid.uuid4()), election_id=election_id, name=jurisdiction_name
     )
-    ja = create_user(user_email)
-    db.session.add(ja)
-    admin = JurisdictionAdministration(user_id=ja.id, jurisdiction_id=jurisdiction.id)
     db.session.add(jurisdiction)
-    db.session.add(admin)
     db.session.commit()
-    return jurisdiction.id, ja.id
+    ja_id = create_jurisdiction_admin(jurisdiction.id, user_email)
+    return jurisdiction.id, ja_id
 
 
 def create_election(