"""Unit tests for the multi-round balanced assignment + completion-code feature. Runs offline (no network / no huggingface_hub calls): assignment.py's core functions are pure, and the round-lifecycle integration tests patch main.py's HF-facing I/O functions with an in-memory fake store. Run with: python -m unittest backend.test_assignment (from repo root) or: python -m unittest test_assignment (from backend/) """ import json import os import random import sys import unittest from datetime import datetime, timedelta, timezone from unittest import mock sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) import assignment # noqa: E402 def make_record(user, video_id, created_at, extra=None): record = { "id": "x", "video_id": video_id, "user": user, "created_at": created_at, "annotations": {"responses": {}}, } if extra: record.update(extra) return record def iso(dt): return dt.isoformat().replace("+00:00", "Z") class DedupeLatestTests(unittest.TestCase): def test_existing_annotations_preserved_and_counted(self): records = [ make_record("Adi", "v1", "2026-01-01T00:00:00Z"), make_record("Adi", "v2", "2026-01-01T00:00:00Z"), make_record("youngsun", "v1", "2026-01-01T00:00:00Z"), ] original_len = len(records) deduped, skipped = assignment.dedupe_latest(records) self.assertEqual(len(records), original_len, "dedupe_latest must not mutate its input") self.assertEqual(skipped, 0) self.assertEqual(len(deduped), 3) by_video = assignment.completed_by_video(deduped) self.assertEqual(by_video["v1"], {"adi", "youngsun"}) self.assertEqual(by_video["v2"], {"adi"}) def test_same_annotator_same_video_counted_once_latest_kept(self): records = [ make_record("Adi", "v1", "2026-01-01T00:00:00Z", {"annotations": {"marker": "old"}}), make_record("Adi", "v1", "2026-01-02T00:00:00Z", {"annotations": {"marker": "new"}}), make_record("adi", "v1", "2026-01-03T00:00:00Z", {"annotations": {"marker": "newest"}}), ] deduped, skipped = assignment.dedupe_latest(records) self.assertEqual(skipped, 0) self.assertEqual(len(deduped), 1, "same annotator + same video, across case variants, is one record") self.assertEqual(deduped[("adi", "v1")]["annotations"]["marker"], "newest") def test_malformed_or_partial_records_skipped_not_counted(self): records = [ make_record("Adi", "v1", "2026-01-01T00:00:00Z"), {"id": "y", "video_id": "v2", "user": "Adi"}, # missing created_at (autosave/reserve-like) {"id": "z", "user": "Adi", "created_at": "2026-01-01T00:00:00Z"}, # missing video_id {"id": "w", "video_id": "v3", "user": "", "created_at": "2026-01-01T00:00:00Z"}, # empty user {"id": "q", "video_id": "v4", "user": "Adi", "created_at": "not-a-date"}, # unparsable timestamp ] deduped, skipped = assignment.dedupe_latest(records) self.assertEqual(len(deduped), 1) self.assertEqual(skipped, 4) class PickBalancedTests(unittest.TestCase): def test_prefers_lowest_coverage(self): coverage = {"v1": 4, "v2": 0, "v3": 2} picked = assignment.pick_balanced(list(coverage), coverage.get, 2, random.Random(0), target=5) self.assertEqual(set(picked), {"v2", "v3"}) def test_excludes_at_or_over_target_when_enough_under_target_remain(self): coverage = {"v1": 5, "v2": 5, "v3": 1} picked = assignment.pick_balanced(list(coverage), coverage.get, 1, random.Random(0), target=5) self.assertEqual(picked, ["v3"]) def test_falls_back_to_at_target_videos_when_not_enough_under_target(self): coverage = {"v1": 5, "v2": 5, "v3": 1} picked = assignment.pick_balanced(list(coverage), coverage.get, 3, random.Random(0), target=5) self.assertEqual(set(picked), {"v1", "v2", "v3"}) def test_never_returns_more_than_requested_or_duplicates(self): coverage = {f"v{i}": i % 3 for i in range(30)} picked = assignment.pick_balanced(list(coverage), coverage.get, 7, random.Random(1), target=5) self.assertEqual(len(picked), 7) self.assertEqual(len(set(picked)), 7) class CompletionCodeTests(unittest.TestCase): def test_format(self): code = assignment.generate_completion_code() self.assertTrue(code.startswith("T2AV-")) suffix = code[len("T2AV-"):] self.assertEqual(len(suffix), 12) self.assertTrue(all(c in assignment.COMPLETION_CODE_ALPHABET for c in suffix)) def test_codes_are_not_trivially_repeated(self): codes = {assignment.generate_completion_code() for _ in range(200)} self.assertEqual(len(codes), 200, "200 draws from a 36^12 space should never collide") class BuildRoundTests(unittest.TestCase): def test_always_excludes_already_completed_never_seeds_with_them(self): """Core policy change: a round is exactly N NEW videos, never padded with history, regardless of how much history exists.""" catalog = [f"v{i}" for i in range(1, 201)] # 200 videos, like production already_completed = set(catalog[:51]) # mirrors real production data (Adi: 51 completed) round_record = assignment.build_round( "adi", "Adi", catalog, already_completed, {}, {}, random.Random(0), round_number=1, ) self.assertEqual(round_record["actual_size"], 20) self.assertEqual(len(round_record["video_ids"]), 20) self.assertEqual( already_completed.intersection(round_record["video_ids"]), set(), "none of the 51 previously completed videos may appear in the new round", ) self.assertEqual(len(set(round_record["video_ids"])), 20, "no duplicates within the round") self.assertEqual(round_record["round_number"], 1) self.assertEqual(round_record["requested_size"], 20) def test_one_historical_completion_gets_twenty_new_not_nineteen(self): catalog = [f"v{i}" for i in range(1, 201)] already_completed = {"v1"} round_record = assignment.build_round( "youngsun", "youngsun_0715", catalog, already_completed, {}, {}, random.Random(0), round_number=1, ) self.assertEqual(round_record["actual_size"], 20) self.assertNotIn("v1", round_record["video_ids"]) def test_brand_new_annotator_gets_exactly_twenty(self): catalog = [f"v{i}" for i in range(1, 201)] round_record = assignment.build_round( "new", "New", catalog, set(), {}, {}, random.Random(0), round_number=1, ) self.assertEqual(round_record["actual_size"], 20) def test_reduced_size_when_not_enough_eligible_remain(self): catalog = [f"v{i}" for i in range(1, 26)] # only 25 videos total already_completed = set(catalog[:15]) # 15 done, only 10 eligible left round_record = assignment.build_round( "x", "X", catalog, already_completed, {}, {}, random.Random(0), round_number=1, ) self.assertEqual(round_record["actual_size"], 10) self.assertEqual(round_record["requested_size"], 20) self.assertEqual(len(round_record["video_ids"]), 10) self.assertEqual(already_completed.intersection(round_record["video_ids"]), set()) def test_second_round_excludes_first_rounds_completions_too(self): """After Adi finishes 51 (history) + 20 (round 1) = 71, round 2 must exclude all 71, selecting from the remaining 129.""" catalog = [f"v{i}" for i in range(1, 201)] historical = set(catalog[:51]) round1 = assignment.build_round( "adi", "Adi", catalog, historical, {}, {}, random.Random(0), round_number=1, ) after_round1 = historical | set(round1["video_ids"]) self.assertEqual(len(after_round1), 71) round2 = assignment.build_round( "adi", "Adi", catalog, after_round1, {}, {}, random.Random(1), round_number=2, ) self.assertEqual(round2["actual_size"], 20) self.assertEqual(after_round1.intersection(round2["video_ids"]), set()) remaining_pool_size = len(catalog) - len(after_round1) self.assertEqual(remaining_pool_size, 129) self.assertEqual(round2["round_number"], 2) def test_existing_completions_affect_new_round_priority(self): catalog = ["v1", "v2", "v3"] completed_by_video_map = {"v1": {"p1", "p2", "p3", "p4"}} # coverage 4, still < target 5 round_record = assignment.build_round( "new", "New", catalog, set(), completed_by_video_map, {}, random.Random(0), round_number=1, videos_per_annotator=2, ) self.assertNotIn("v1", round_record["video_ids"], "v2/v3 have lower coverage and must be preferred") self.assertEqual(set(round_record["video_ids"]), {"v2", "v3"}) class ReservedByVideoTests(unittest.TestCase): def test_counts_open_uncompleted_unexpired_reservations(self): now = datetime(2026, 1, 10, tzinfo=timezone.utc) assignments = [ {"annotator_id": "a", "video_ids": ["v1", "v2"], "created_at": "2026-01-09T00:00:00Z"}, {"annotator_id": "b", "video_ids": ["v1"], "created_at": "2026-01-09T00:00:00Z"}, ] completed_by_annotator = {"a": {"v2"}} reserved = assignment.reserved_by_video(assignments, completed_by_annotator, now, ttl_seconds=7 * 24 * 3600) self.assertEqual(reserved.get("v1"), 2) self.assertNotIn("v2", reserved) def test_excludes_expired_reservations_without_needing_to_delete_anything(self): now = datetime(2026, 1, 10, tzinfo=timezone.utc) assignments = [{"annotator_id": "a", "video_ids": ["v1"], "created_at": "2025-01-01T00:00:00Z"}] reserved = assignment.reserved_by_video(assignments, {}, now, ttl_seconds=24 * 3600) self.assertEqual(reserved.get("v1", 0), 0) def test_excludes_the_requesting_annotator_from_their_own_reservations(self): now = datetime(2026, 1, 10, tzinfo=timezone.utc) assignments = [{"annotator_id": "a", "video_ids": ["v1"], "created_at": "2026-01-09T00:00:00Z"}] reserved = assignment.reserved_by_video(assignments, {}, now, ttl_seconds=24 * 3600, exclude_annotator="a") self.assertEqual(reserved.get("v1", 0), 0) class SequentialBalanceTests(unittest.TestCase): def test_sequential_rounds_stay_reasonably_balanced(self): catalog = [f"v{i}" for i in range(1, 81)] # 80 videos reserved_by_video_map = {} rng = random.Random(42) for i in range(10): # 10 annotators * 20 slots = 200 slots over 80 videos round_record = assignment.build_round( f"person{i}", f"person{i}", catalog, set(), {}, reserved_by_video_map, rng, round_number=1, ) for video_id in round_record["video_ids"]: reserved_by_video_map[video_id] = reserved_by_video_map.get(video_id, 0) + 1 counts = list(reserved_by_video_map.values()) self.assertEqual(len(reserved_by_video_map), 80, "every video should have been touched") self.assertLessEqual(max(counts) - min(counts), 2, "balancing should keep coverage tight") class FakeRepoStore: """Stands in for the annotations dataset repo's assignments/ and completions/ prefixes.""" def __init__(self): self.files = {} def read(self, path): return self.files.get(path) def write_kwargs(self, **kwargs): path = kwargs["path_in_repo"] content = json.loads(kwargs["path_or_fileobj"].getvalue().decode("utf-8")) self.files[path] = content def list_active_round_records(self): records = [] for path, pointer in self.files.items(): if not path.endswith("/current.json"): continue annotator_id = path.split("/")[1] round_path = f"assignments/{annotator_id}/rounds/{pointer['assignment_id']}.json" completion_path = f"completions/{annotator_id}/{pointer['assignment_id']}.json" round_record = self.files.get(round_path) if round_record is None or completion_path in self.files: continue records.append(round_record) return records class RoundLifecycleIntegrationTests(unittest.TestCase): """Exercises main.get_or_create_current_round() with every HF network call replaced by an in-memory fake, so this stays fully offline.""" def setUp(self): import main as backend_main self.backend_main = backend_main self.store = FakeRepoStore() self.catalog = [f"v{i}" for i in range(1, 201)] # 200 videos, like production self.completed_records = [] # raw annotation records fed to both list_annotation_records variants patches = [ mock.patch.object(backend_main, "_read_json_from_repo", side_effect=self.store.read), # Both the cached and fresh variants must be mocked - the completion # check in get_or_create_current_round() deliberately bypasses the # cache (fresh=True) so a real annotator sees their code promptly; # leaving list_annotation_records() unmocked would silently fall # through to a real network call here. mock.patch.object( backend_main, "list_annotation_records_cached", side_effect=lambda: self.completed_records ), mock.patch.object( backend_main, "list_annotation_records", side_effect=lambda: self.completed_records ), mock.patch.object( backend_main, "list_active_round_records", side_effect=self.store.list_active_round_records ), mock.patch.object(backend_main, "ensure_dataset_exists", return_value=None), mock.patch.object(backend_main, "ensure_dataset_configured", return_value=None), mock.patch.object(backend_main, "_current_catalog_video_ids", return_value=self.catalog), mock.patch.object(backend_main, "upload_with_retry", side_effect=self.store.write_kwargs), ] for patcher in patches: patcher.start() self.addCleanup(patcher.stop) def _complete_round(self, annotator_raw, round_record): """Simulates the annotator finishing every video in a round by appending matching annotation records - exactly what save_annotation would have produced.""" now = datetime.now(timezone.utc) for video_id in round_record["video_ids"]: self.completed_records.append(make_record(annotator_raw, video_id, iso(now))) def test_brand_new_annotator_gets_round_one_of_twenty(self): result = self.backend_main.get_or_create_current_round("Scratch Tester") self.assertEqual(result["status"], "in_progress") self.assertEqual(result["round_number"], 1) self.assertEqual(len(result["video_ids"]), 20) self.assertEqual(len(set(result["video_ids"])), 20) def test_refresh_while_incomplete_returns_identical_round(self): first = self.backend_main.get_or_create_current_round("Scratch Tester") second = self.backend_main.get_or_create_current_round("Scratch Tester") self.assertEqual(first["video_ids"], second["video_ids"]) self.assertEqual(first["assignment_id"], second["assignment_id"]) def test_no_completion_code_at_nineteen_of_twenty(self): first = self.backend_main.get_or_create_current_round("Scratch Tester") for video_id in first["video_ids"][:19]: self.completed_records.append(make_record("Scratch Tester", video_id, iso(datetime.now(timezone.utc)))) result = self.backend_main.get_or_create_current_round("Scratch Tester") self.assertEqual(result["status"], "in_progress") self.assertNotIn("completion_code", result) def test_completion_code_persisted_once_at_twenty_of_twenty(self): first = self.backend_main.get_or_create_current_round("Scratch Tester") self._complete_round("Scratch Tester", first) result_a = self.backend_main.get_or_create_current_round("Scratch Tester") result_b = self.backend_main.get_or_create_current_round("Scratch Tester") self.assertEqual(result_a["status"], "completed") self.assertTrue(result_a["completion_code"].startswith("T2AV-")) self.assertEqual(result_a["completion_code"], result_b["completion_code"], "code must be stable on refresh") completion_files = [p for p in self.store.files if p.startswith("completions/")] self.assertEqual(len(completion_files), 1, "completion record written exactly once") def test_within_grace_period_shows_same_completed_round_not_a_new_one(self): first = self.backend_main.get_or_create_current_round("Scratch Tester") self._complete_round("Scratch Tester", first) self.backend_main.get_or_create_current_round("Scratch Tester") # triggers completion-record creation # Directly age the completion record to just inside the grace window. completion_path = f"completions/scratch-tester/{first['assignment_id']}.json" recent = datetime.now(timezone.utc) - timedelta( seconds=assignment.COMPLETION_GRACE_PERIOD_SECONDS - 30 ) self.store.files[completion_path]["completed_at"] = iso(recent) result = self.backend_main.get_or_create_current_round("Scratch Tester") self.assertEqual(result["status"], "completed") self.assertEqual(result["round_number"], 1) def test_past_grace_period_auto_creates_next_round_with_no_extra_call(self): first = self.backend_main.get_or_create_current_round("Scratch Tester") self._complete_round("Scratch Tester", first) self.backend_main.get_or_create_current_round("Scratch Tester") # triggers completion-record creation # Directly age the completion record past the grace window - simulates "came back later". completion_path = f"completions/scratch-tester/{first['assignment_id']}.json" stale = datetime.now(timezone.utc) - timedelta( seconds=assignment.COMPLETION_GRACE_PERIOD_SECONDS + 60 ) self.store.files[completion_path]["completed_at"] = iso(stale) result = self.backend_main.get_or_create_current_round("Scratch Tester") self.assertEqual(result["status"], "in_progress") self.assertEqual(result["round_number"], 2) self.assertEqual(set(result["video_ids"]).intersection(first["video_ids"]), set()) def test_never_assigns_the_same_video_twice_across_rounds(self): first = self.backend_main.get_or_create_current_round("Scratch Tester") self._complete_round("Scratch Tester", first) self.backend_main.get_or_create_current_round("Scratch Tester") completion_path = f"completions/scratch-tester/{first['assignment_id']}.json" stale = datetime.now(timezone.utc) - timedelta(seconds=assignment.COMPLETION_GRACE_PERIOD_SECONDS + 60) self.store.files[completion_path]["completed_at"] = iso(stale) second = self.backend_main.get_or_create_current_round("Scratch Tester") all_seen = first["video_ids"] + second["video_ids"] self.assertEqual(len(all_seen), len(set(all_seen))) def test_different_annotators_each_get_their_own_round(self): a = self.backend_main.get_or_create_current_round("Person A") b = self.backend_main.get_or_create_current_round("Person B") self.assertEqual(len(a["video_ids"]), 20) self.assertEqual(len(b["video_ids"]), 20) round_files = [p for p in self.store.files if "/rounds/" in p] self.assertEqual(len(round_files), 2) def test_simultaneous_requests_do_not_create_duplicate_rounds(self): """The lock fully serializes get_or_create_current_round, so we simulate "simultaneous" by calling it back-to-back and asserting only one round/pointer combination was ever written - the real concurrency guarantee is exercised live against a scratch HF repo separately.""" for _ in range(5): self.backend_main.get_or_create_current_round("Scratch Tester") round_files = [p for p in self.store.files if "/rounds/" in p] self.assertEqual(len(round_files), 1) def test_partial_records_never_mistakenly_counted_as_completed(self): first = self.backend_main.get_or_create_current_round("Scratch Tester") # An autosave/reservation-like record missing created_at for every video. for video_id in first["video_ids"]: self.completed_records.append({"id": "x", "video_id": video_id, "user": "Scratch Tester"}) result = self.backend_main.get_or_create_current_round("Scratch Tester") self.assertEqual(result["status"], "in_progress") if __name__ == "__main__": unittest.main()