From e2d6e1bf098671c9b57cf24f1d08fbf833b4164b Mon Sep 17 00:00:00 2001 From: Nikhil Bhatia Date: Tue, 16 Jun 2026 16:17:43 +0000 Subject: [PATCH 1/3] Batch comparison: Add extra batches once per round --- .basedpyright/baseline.json | 288 ++++++++++++++++++++++++++++++++++-- server/api/shared.py | 88 +++++++---- 2 files changed, 339 insertions(+), 37 deletions(-) diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 28c002000..04be47819 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -36953,6 +36953,70 @@ "lineCount": 1 } }, + { + "code": "reportUnknownVariableType", + "range": { + "startColumn": 4, + "endColumn": 17, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 8, + "endColumn": 81, + "lineCount": 1 + } + }, + { + "code": "reportUnknownVariableType", + "range": { + "startColumn": 54, + "endColumn": 55, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 59, + "endColumn": 80, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 12, + "endColumn": 28, + "lineCount": 1 + } + }, + { + "code": "reportUnknownLambdaType", + "range": { + "startColumn": 19, + "endColumn": 20, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 22, + "endColumn": 28, + "lineCount": 1 + } + }, + { + "code": "reportUnknownLambdaType", + "range": { + "startColumn": 22, + "endColumn": 28, + "lineCount": 1 + } + }, { "code": "reportUnknownVariableType", "range": { @@ -36969,6 +37033,158 @@ "lineCount": 1 } }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 12, + "endColumn": 52, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 53, + "endColumn": 68, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 53, + "endColumn": 68, + "lineCount": 1 + } + }, + { + "code": "reportUnknownVariableType", + "range": { + "startColumn": 4, + "endColumn": 14, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 8, + "endColumn": 24, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 8, + "endColumn": 15, + "lineCount": 2 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 8, + "endColumn": 22, + "lineCount": 3 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 8, + "endColumn": 12, + "lineCount": 4 + } + }, + { + "code": "reportUnknownVariableType", + "range": { + "startColumn": 4, + "endColumn": 19, + "lineCount": 1 + } + }, + { + "code": "reportUnknownVariableType", + "range": { + "startColumn": 12, + "endColumn": 29, + "lineCount": 1 + } + }, + { + "code": "reportUnknownVariableType", + "range": { + "startColumn": 31, + "endColumn": 41, + "lineCount": 1 + } + }, + { + "code": "reportUnknownVariableType", + "range": { + "startColumn": 43, + "endColumn": 51, + "lineCount": 1 + } + }, + { + "code": "reportUnknownVariableType", + "range": { + "startColumn": 4, + "endColumn": 19, + "lineCount": 1 + } + }, + { + "code": "reportUnknownVariableType", + "range": { + "startColumn": 12, + "endColumn": 29, + "lineCount": 1 + } + }, + { + "code": "reportUnknownVariableType", + "range": { + "startColumn": 31, + "endColumn": 41, + "lineCount": 1 + } + }, + { + "code": "reportUnknownVariableType", + "range": { + "startColumn": 43, + "endColumn": 51, + "lineCount": 1 + } + }, + { + "code": "reportUnknownVariableType", + "range": { + "startColumn": 12, + "endColumn": 24, + "lineCount": 1 + } + }, + { + "code": "reportUnknownVariableType", + "range": { + "startColumn": 12, + "endColumn": 37, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 70, + "endColumn": 85, + "lineCount": 1 + } + }, { "code": "reportUnknownVariableType", "range": { @@ -37114,18 +37330,18 @@ } }, { - "code": "reportAny", + "code": "reportUnknownVariableType", "range": { - "startColumn": 20, - "endColumn": 21, + "startColumn": 12, + "endColumn": 30, "lineCount": 1 } }, { "code": "reportUnknownMemberType", "range": { - "startColumn": 35, - "endColumn": 52, + "startColumn": 60, + "endColumn": 77, "lineCount": 1 } }, @@ -37193,6 +37409,14 @@ "lineCount": 1 } }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 35, + "endColumn": 60, + "lineCount": 1 + } + }, { "code": "reportUnknownArgumentType", "range": { @@ -37249,6 +37473,14 @@ "lineCount": 1 } }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 35, + "endColumn": 60, + "lineCount": 1 + } + }, { "code": "reportUnknownParameterType", "range": { @@ -37393,6 +37625,14 @@ "lineCount": 1 } }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 35, + "endColumn": 60, + "lineCount": 1 + } + }, { "code": "reportUnknownVariableType", "range": { @@ -37401,27 +37641,43 @@ "lineCount": 1 } }, + { + "code": "reportUnknownVariableType", + "range": { + "startColumn": 12, + "endColumn": 37, + "lineCount": 1 + } + }, { "code": "reportUnknownMemberType", "range": { - "startColumn": 28, - "endColumn": 49, + "startColumn": 70, + "endColumn": 85, "lineCount": 1 } }, { - "code": "reportAny", + "code": "reportUnknownVariableType", "range": { - "startColumn": 20, - "endColumn": 21, + "startColumn": 12, + "endColumn": 48, "lineCount": 1 } }, { "code": "reportUnknownMemberType", "range": { - "startColumn": 35, - "endColumn": 52, + "startColumn": 60, + "endColumn": 77, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 19, + "endColumn": 55, "lineCount": 1 } }, @@ -37505,6 +37761,14 @@ "lineCount": 1 } }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 39, + "endColumn": 64, + "lineCount": 1 + } + }, { "code": "reportUnknownVariableType", "range": { diff --git a/server/api/shared.py b/server/api/shared.py index 9b9dec62a..7764a05ba 100644 --- a/server/api/shared.py +++ b/server/api/shared.py @@ -635,7 +635,6 @@ class BatchDraw(TypedDict): def compute_sample_batches_for_contest( election: Election, - round_num: int, contest: Contest, contest_sample_size: SampleSize, ) -> list[BatchDraw]: @@ -681,12 +680,49 @@ def compute_sample_batches_for_contest( for ticket_number, batch_key in sample ] - # Experimental feature - # Add extra batches on top of the original sample that will be audited, but - # not counted in the final risk measurement. + return sample_batches + + +# Experimental feature +# Add extra batches on top of the original sample that will be audited, but +# not counted in the final risk measurement. +def compute_extra_batches_for_round( + election: Election, + round_num: int, + contest_sample_sizes: list[tuple[Contest, SampleSize]], + sampled_batches: list[BatchDraw], +) -> list[BatchDraw]: + jurisdictions = sorted( + {j for contest, _ in contest_sample_sizes for j in contest.jurisdictions}, + key=lambda j: j.name, + ) + # Each unique jurisdiction mapped to the first audited contest it appears in. + jurisdiction_id_to_contest_id = {} + for contest, _ in contest_sample_sizes: + for jurisdiction in contest.jurisdictions: + jurisdiction_id_to_contest_id.setdefault(jurisdiction.id, contest.id) + + batch_rows = ( + Batch.query.join(Jurisdiction) + .filter(Jurisdiction.election_id == election.id) + .with_entities(Jurisdiction.name, Batch.name, Batch.id) + .all() + ) + batch_key_to_id = { + (jurisdiction_name, batch_name): batch_id + for jurisdiction_name, batch_name, batch_id in batch_rows + } + batch_id_to_key = { + batch_id: (jurisdiction_name, batch_name) + for jurisdiction_name, batch_name, batch_id in batch_rows + } + + extra_batches: list[BatchDraw] = [] + if is_enabled_sample_extra_batches_by_counting_group(election) and round_num == 1: rand = random.Random(str(election.random_seed)) - for jurisdiction in contest.jurisdictions: + for jurisdiction in jurisdictions: + representative_contest_id = jurisdiction_id_to_contest_id[jurisdiction.id] batch_key_to_num_ballots = { (jurisdiction.name, batch.name): batch.num_ballots for batch in jurisdiction.batches @@ -711,9 +747,9 @@ def compute_sample_batches_for_contest( in [CountingGroup.ABSENTEE_BY_MAIL, CountingGroup.PROVISIONAL] } sampled_batch_keys = { - batch_key - for _, batch_key in sample - if batch_key[0] == jurisdiction.name + batch_id_to_key[batch["batch_id"]] + for batch in sampled_batches + if batch_id_to_key[batch["batch_id"]][0] == jurisdiction.name } extra_batch_keys = set() @@ -724,10 +760,10 @@ def compute_sample_batches_for_contest( ): extra_bmd_batch_key = rand.choice(sorted(bmd_batch_keys)) extra_batch_keys.add(extra_bmd_batch_key) - sample_batches.append( + extra_batches.append( BatchDraw( batch_id=batch_key_to_id[extra_bmd_batch_key], - contest_id=contest.id, + contest_id=representative_contest_id, ticket_number=EXTRA_TICKET_NUMBER, ) ) @@ -738,10 +774,10 @@ def compute_sample_batches_for_contest( ): extra_hmpb_batch_key = rand.choice(sorted(hmpb_batch_keys)) extra_batch_keys.add(extra_hmpb_batch_key) - sample_batches.append( + extra_batches.append( BatchDraw( batch_id=batch_key_to_id[extra_hmpb_batch_key], - contest_id=contest.id, + contest_id=representative_contest_id, ticket_number=EXTRA_TICKET_NUMBER, ) ) @@ -773,21 +809,22 @@ def compute_percentage_of_jurisdiction_ballots_selected( ) extra_batch_key = rand.choice(sorted(remaining_batch_keys)) extra_batch_keys.add(extra_batch_key) - sample_batches.append( + extra_batches.append( BatchDraw( batch_id=batch_key_to_id[extra_batch_key], - contest_id=contest.id, + contest_id=representative_contest_id, ticket_number=EXTRA_TICKET_NUMBER, ) ) if is_enabled_sample_extra_batches_to_ensure_one_per_jurisdiction(election): rand = random.Random(str(election.random_seed)) - for jurisdiction in contest.jurisdictions: + for jurisdiction in jurisdictions: + representative_contest_id = jurisdiction_id_to_contest_id[jurisdiction.id] sampled_batch_keys_from_jurisdiction = { - batch_key - for _, batch_key in sample - if batch_key[0] == jurisdiction.name + batch_id_to_key[batch["batch_id"]] + for batch in sampled_batches + if batch_id_to_key[batch["batch_id"]][0] == jurisdiction.name } # If we didn't sample any batches from this jurisdiction, add one if len(sampled_batch_keys_from_jurisdiction) == 0: @@ -796,15 +833,15 @@ def compute_percentage_of_jurisdiction_ballots_selected( } if len(jurisdiction_batch_keys) > 0: extra_batch_key = rand.choice(sorted(jurisdiction_batch_keys)) - sample_batches.append( + extra_batches.append( BatchDraw( batch_id=batch_key_to_id[extra_batch_key], - contest_id=contest.id, + contest_id=representative_contest_id, ticket_number=EXTRA_TICKET_NUMBER, ) ) - return sample_batches + return extra_batches def compute_sample_batches( @@ -815,11 +852,12 @@ def compute_sample_batches( sample_batches = [ batch for contest, sample_size in contest_sample_sizes - for batch in compute_sample_batches_for_contest( - election, round_num, contest, sample_size - ) + for batch in compute_sample_batches_for_contest(election, contest, sample_size) ] - return sample_batches + extra_batches = compute_extra_batches_for_round( + election, round_num, contest_sample_sizes, sample_batches + ) + return sample_batches + extra_batches def compute_sample_ballots( From fd50a7cbde48057cb3adc231ed6a4ac8750c160f Mon Sep 17 00:00:00 2001 From: Nikhil Bhatia Date: Tue, 16 Jun 2026 22:20:38 +0000 Subject: [PATCH 2/3] Batch comparison: Exclude all extra batches from sampled_batch_rows() when flag is set --- .basedpyright/baseline.json | 76 ++++++++++++++++++++++++++++++++++++- server/api/shared.py | 14 +++++-- 2 files changed, 85 insertions(+), 5 deletions(-) diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 04be47819..bd22dc6c1 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -35142,9 +35142,57 @@ "range": { "startColumn": 12, "endColumn": 17, + "lineCount": 2 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 12, + "endColumn": 17, + "lineCount": 3 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 12, + "endColumn": 19, "lineCount": 4 } }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 12, + "endColumn": 19, + "lineCount": 5 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 12, + "endColumn": 50, + "lineCount": 5 + } + }, + { + "code": "reportUnknownVariableType", + "range": { + "startColumn": 8, + "endColumn": 30, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 12, + "endColumn": 41, + "lineCount": 1 + } + }, { "code": "reportUnknownMemberType", "range": { @@ -35157,16 +35205,40 @@ "code": "reportUnknownMemberType", "range": { "startColumn": 12, - "endColumn": 19, + "endColumn": 17, "lineCount": 6 } }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 12, + "endColumn": 19, + "lineCount": 7 + } + }, { "code": "reportUnknownArgumentType", "range": { "startColumn": 12, "endColumn": 50, - "lineCount": 6 + "lineCount": 7 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 16, + "endColumn": 57, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 16, + "endColumn": 69, + "lineCount": 1 } }, { diff --git a/server/api/shared.py b/server/api/shared.py index 7764a05ba..7e3c67d5d 100644 --- a/server/api/shared.py +++ b/server/api/shared.py @@ -217,8 +217,16 @@ def sampled_batch_results( # and audit reports if not include_non_rla_batches: extra_batch_keys = set( - SampledBatchDraw.query.filter_by( - contest_id=contest.id, ticket_number=EXTRA_TICKET_NUMBER + SampledBatchDraw.query.filter_by(ticket_number=EXTRA_TICKET_NUMBER) + .join(Batch) + .join(Jurisdiction) + .filter(Jurisdiction.election_id == contest.election_id) + .values(Jurisdiction.name, Batch.name) + ) + rla_sampled_batch_keys = set( + SampledBatchDraw.query.filter( + SampledBatchDraw.contest_id == contest.id, + SampledBatchDraw.ticket_number != EXTRA_TICKET_NUMBER, ) .join(Batch) .join(Jurisdiction) @@ -227,7 +235,7 @@ def sampled_batch_results( results = { batch_key: result for batch_key, result in results.items() - if batch_key not in extra_batch_keys + if batch_key not in (extra_batch_keys - rla_sampled_batch_keys) } return results From 9d55661f11190423b0110b494854a80eeee62df2 Mon Sep 17 00:00:00 2001 From: Nikhil Bhatia Date: Tue, 16 Jun 2026 23:28:22 +0000 Subject: [PATCH 3/3] Add test for multi-contest with extra batches by counting group --- .basedpyright/baseline.json | 264 ++++++++++++++++++ ..._sample_extra_batches_by_counting_group.py | 126 +++++++++ 2 files changed, 390 insertions(+) diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index bd22dc6c1..7f8b2e4dd 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -118982,6 +118982,270 @@ "endColumn": 41, "lineCount": 1 } + }, + { + "code": "reportUnusedParameter", + "range": { + "startColumn": 4, + "endColumn": 10, + "lineCount": 1 + } + }, + { + "code": "reportUnknownParameterType", + "range": { + "startColumn": 4, + "endColumn": 21, + "lineCount": 1 + } + }, + { + "code": "reportMissingParameterType", + "range": { + "startColumn": 4, + "endColumn": 21, + "lineCount": 1 + } + }, + { + "code": "reportUnusedParameter", + "range": { + "startColumn": 4, + "endColumn": 21, + "lineCount": 1 + } + }, + { + "code": "reportUnknownParameterType", + "range": { + "startColumn": 4, + "endColumn": 13, + "lineCount": 1 + } + }, + { + "code": "reportMissingParameterType", + "range": { + "startColumn": 4, + "endColumn": 13, + "lineCount": 1 + } + }, + { + "code": "reportUnusedParameter", + "range": { + "startColumn": 4, + "endColumn": 13, + "lineCount": 1 + } + }, + { + "code": "reportAny", + "range": { + "startColumn": 4, + "endColumn": 6, + "lineCount": 1 + } + }, + { + "code": "reportAny", + "range": { + "startColumn": 14, + "endColumn": 16, + "lineCount": 1 + } + }, + { + "code": "reportImplicitStringConcatenation", + "range": { + "startColumn": 12, + "endColumn": 41, + "lineCount": 11 + } + }, + { + "code": "reportImplicitStringConcatenation", + "range": { + "startColumn": 12, + "endColumn": 36, + "lineCount": 8 + } + }, + { + "code": "reportAny", + "range": { + "startColumn": 4, + "endColumn": 23, + "lineCount": 1 + } + }, + { + "code": "reportAny", + "range": { + "startColumn": 37, + "endColumn": 44, + "lineCount": 1 + } + }, + { + "code": "reportAny", + "range": { + "startColumn": 37, + "endColumn": 44, + "lineCount": 1 + } + }, + { + "code": "reportAny", + "range": { + "startColumn": 4, + "endColumn": 6, + "lineCount": 1 + } + }, + { + "code": "reportAny", + "range": { + "startColumn": 20, + "endColumn": 30, + "lineCount": 1 + } + }, + { + "code": "reportAny", + "range": { + "startColumn": 32, + "endColumn": 39, + "lineCount": 1 + } + }, + { + "code": "reportAny", + "range": { + "startColumn": 43, + "endColumn": 68, + "lineCount": 1 + } + }, + { + "code": "reportAny", + "range": { + "startColumn": 14, + "endColumn": 16, + "lineCount": 1 + } + }, + { + "code": "reportAny", + "range": { + "startColumn": 4, + "endColumn": 14, + "lineCount": 1 + } + }, + { + "code": "reportAny", + "range": { + "startColumn": 28, + "endColumn": 35, + "lineCount": 1 + } + }, + { + "code": "reportAny", + "range": { + "startColumn": 28, + "endColumn": 35, + "lineCount": 1 + } + }, + { + "code": "reportUnknownVariableType", + "range": { + "startColumn": 4, + "endColumn": 15, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 8, + "endColumn": 35, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 8, + "endColumn": 15, + "lineCount": 2 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 8, + "endColumn": 12, + "lineCount": 7 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 15, + "endColumn": 26, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 13, + "endColumn": 26, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 28, + "endColumn": 43, + "lineCount": 1 + } + }, + { + "code": "reportUnknownVariableType", + "range": { + "startColumn": 49, + "endColumn": 53, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 15, + "endColumn": 56, + "lineCount": 1 + } + }, + { + "code": "reportUnknownMemberType", + "range": { + "startColumn": 16, + "endColumn": 31, + "lineCount": 1 + } + }, + { + "code": "reportUnknownVariableType", + "range": { + "startColumn": 36, + "endColumn": 40, + "lineCount": 1 + } } ], "./server/tests/batch_comparison/test_sample_extra_batches_to_ensure_one_per_jurisdiction.py": [ diff --git a/server/tests/batch_comparison/test_sample_extra_batches_by_counting_group.py b/server/tests/batch_comparison/test_sample_extra_batches_by_counting_group.py index 33434e09e..451eb7a2c 100644 --- a/server/tests/batch_comparison/test_sample_extra_batches_by_counting_group.py +++ b/server/tests/batch_comparison/test_sample_extra_batches_by_counting_group.py @@ -763,3 +763,129 @@ def test_sample_extra_batches_with_combined_batches( # Check the audit report rv = client.get(f"/api/election/{election_id}/report") assert_match_report(rv.data, snapshot) + + +@pytest.mark.parametrize( + "org_id", + ["TEST-ORG/sample-extra-batches-by-counting-group"], + indirect=True, +) +def test_sample_extra_batches_multi_contest( + client: FlaskClient, + org_id: str, + election_id: str, + jurisdiction_ids: list[str], + election_settings, + manifests, +): + # Two targeted contests where jurisdiction_ids[0] participates in both and + # jurisdiction_ids[1] is exclusive to Contest 1 + set_logged_in_user(client, UserType.AUDIT_ADMIN, DEFAULT_AA_EMAIL) + contests = [ + { + "id": str(uuid.uuid4()), + "name": "Contest 1", + "isTargeted": True, + "choices": [ + {"id": str(uuid.uuid4()), "name": "candidate 1", "numVotes": 5000}, + {"id": str(uuid.uuid4()), "name": "candidate 2", "numVotes": 2500}, + {"id": str(uuid.uuid4()), "name": "candidate 3", "numVotes": 2500}, + ], + "numWinners": 1, + "votesAllowed": 2, + "jurisdictionIds": jurisdiction_ids[:2], + }, + { + "id": str(uuid.uuid4()), + "name": "Contest 2", + "isTargeted": True, + "choices": [ + {"id": str(uuid.uuid4()), "name": "candidate 4", "numVotes": 2500}, + {"id": str(uuid.uuid4()), "name": "candidate 5", "numVotes": 1250}, + ], + "numWinners": 1, + "votesAllowed": 2, + "jurisdictionIds": [jurisdiction_ids[0]], + }, + ] + rv = put_json(client, f"/api/election/{election_id}/contest", contests) + assert_ok(rv) + + # Batch tallies. jurisdiction_ids[0] is in both contests; + # jurisdiction_ids[1] is in Contest 1 only. + set_logged_in_user( + client, UserType.JURISDICTION_ADMIN, default_ja_email(election_id) + ) + rv = upload_batch_tallies( + client, + io.BytesIO( + b"Batch Name,Contest 1 - candidate 1,Contest 1 - candidate 2," + b"Contest 1 - candidate 3,Contest 2 - candidate 4,Contest 2 - candidate 5\n" + b"Batch 1,500,250,250,500,250\n" + b"Batch 2,500,250,250,500,250\n" + b"Batch 3,500,250,250,500,250\n" + b"Batch 4,500,250,250,500,250\n" + b"Batch 5,100,50,50,100,50\n" + b"Batch 6,100,50,50,100,50\n" + b"Batch 7,100,50,50,100,50\n" + b"Batch 8,100,50,50,100,50\n" + b"Batch 9,100,50,50,100,50\n" + ), + election_id, + jurisdiction_ids[0], + ) + assert_ok(rv) + rv = upload_batch_tallies( + client, + io.BytesIO( + b"Batch Name,Contest 1 - candidate 1,Contest 1 - candidate 2," + b"Contest 1 - candidate 3\n" + b"Batch 1,500,250,250\n" + b"Batch 2,500,250,250\n" + b"Batch 3,500,250,250\n" + b"Batch 4,500,250,250\n" + b"Batch 5,300,100,100\n" + b"Batch 6,200,150,150\n" + ), + election_id, + jurisdiction_ids[1], + ) + assert_ok(rv) + + # Start round 1, which triggers extra-batch selection + set_logged_in_user(client, UserType.AUDIT_ADMIN, DEFAULT_AA_EMAIL) + rv = client.get(f"/api/election/{election_id}/sample-sizes/1") + assert rv.status_code == 200 + sample_size_options = json.loads(rv.data)["sampleSizes"] + rv = post_json( + client, + f"/api/election/{election_id}/round", + { + "roundNum": 1, + "sampleSizes": { + contest_id: options[0] + for contest_id, options in sample_size_options.items() + }, + }, + ) + assert_ok(rv) + rv = client.get(f"/api/election/{election_id}/round") + round_1_id = json.loads(rv.data)["rounds"][0]["id"] + + # jurisdiction_ids[0]'s regular sample covers either its HMPB batch or its + # BMD batch, so it needs one extra batch for counting-group coverage added + # for the round + extra_draws = ( + SampledBatchDraw.query.join(Batch) + .filter( + Batch.jurisdiction_id == jurisdiction_ids[0], + SampledBatchDraw.round_id == round_1_id, + SampledBatchDraw.ticket_number == EXTRA_TICKET_NUMBER, + ) + .all() + ) + assert len(extra_draws) == 1, ( + "expected one extra batch for the multi-contest jurisdiction, got " + f"{[(draw.batch_id, draw.contest_id) for draw in extra_draws]}" + ) + assert len({draw.contest_id for draw in extra_draws}) == 1