dghadiya/t2av_eval

1 2"""Unit tests for the multi-round balanced assignment + completion-code3feature.4 5Runs offline (no network / no huggingface_hub calls): assignment.py's core6functions are pure, and the round-lifecycle integration tests patch main.py's7HF-facing I/O functions with an in-memory fake store.8 9Run with:  python -m unittest backend.test_assignment   (from repo root)10       or: python -m unittest test_assignment            (from backend/)11"""12 13import json14import os15import random16import sys17import unittest18from datetime import datetime, timedelta, timezone19from unittest import mock20 21sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))22 23import assignment  # noqa: E40224 25 26def make_record(user, video_id, created_at, extra=None):27    record = {28        "id": "x",29        "video_id": video_id,30        "user": user,31        "created_at": created_at,32        "annotations": {"responses": {}},33    }34    if extra:35        record.update(extra)36    return record37 38 39def iso(dt):40    return dt.isoformat().replace("+00:00", "Z")41 42 43class DedupeLatestTests(unittest.TestCase):44    def test_existing_annotations_preserved_and_counted(self):45        records = [46            make_record("Adi", "v1", "2026-01-01T00:00:00Z"),47            make_record("Adi", "v2", "2026-01-01T00:00:00Z"),48            make_record("youngsun", "v1", "2026-01-01T00:00:00Z"),49        ]50        original_len = len(records)51 52        deduped, skipped = assignment.dedupe_latest(records)53 54        self.assertEqual(len(records), original_len, "dedupe_latest must not mutate its input")55        self.assertEqual(skipped, 0)56        self.assertEqual(len(deduped), 3)57        by_video = assignment.completed_by_video(deduped)58        self.assertEqual(by_video["v1"], {"adi", "youngsun"})59        self.assertEqual(by_video["v2"], {"adi"})60 61    def test_same_annotator_same_video_counted_once_latest_kept(self):62        records = [63            make_record("Adi", "v1", "2026-01-01T00:00:00Z", {"annotations": {"marker": "old"}}),64            make_record("Adi", "v1", "2026-01-02T00:00:00Z", {"annotations": {"marker": "new"}}),65            make_record("adi", "v1", "2026-01-03T00:00:00Z", {"annotations": {"marker": "newest"}}),66        ]67 68        deduped, skipped = assignment.dedupe_latest(records)69 70        self.assertEqual(skipped, 0)71        self.assertEqual(len(deduped), 1, "same annotator + same video, across case variants, is one record")72        self.assertEqual(deduped[("adi", "v1")]["annotations"]["marker"], "newest")73 74    def test_malformed_or_partial_records_skipped_not_counted(self):75        records = [76            make_record("Adi", "v1", "2026-01-01T00:00:00Z"),77            {"id": "y", "video_id": "v2", "user": "Adi"},  # missing created_at (autosave/reserve-like)78            {"id": "z", "user": "Adi", "created_at": "2026-01-01T00:00:00Z"},  # missing video_id79            {"id": "w", "video_id": "v3", "user": "", "created_at": "2026-01-01T00:00:00Z"},  # empty user80            {"id": "q", "video_id": "v4", "user": "Adi", "created_at": "not-a-date"},  # unparsable timestamp81        ]82 83        deduped, skipped = assignment.dedupe_latest(records)84 85        self.assertEqual(len(deduped), 1)86        self.assertEqual(skipped, 4)87 88 89class SavedResponsesForRoundTests(unittest.TestCase):90    def test_returns_responses_only_for_this_annotator_and_this_rounds_videos(self):91        records = [92            make_record("Adi", "v1", "2026-01-01T00:00:00Z", {"annotations": {"responses": {"marker": "adi-v1"}}}),93            make_record("Adi", "v2", "2026-01-01T00:00:00Z", {"annotations": {"responses": {"marker": "adi-v2"}}}),94            make_record(95                "youngsun", "v1", "2026-01-01T00:00:00Z", {"annotations": {"responses": {"marker": "youngsun-v1"}}}96            ),97        ]98        deduped, _skipped = assignment.dedupe_latest(records)99 100        result = assignment.saved_responses_for_round(deduped, "adi", ["v1", "v2", "v3"])101 102        self.assertEqual(result, {"v1": {"marker": "adi-v1"}, "v2": {"marker": "adi-v2"}})103 104    def test_excludes_videos_outside_the_given_round(self):105        records = [106            make_record("Adi", "v9", "2026-01-01T00:00:00Z", {"annotations": {"responses": {"marker": "old-round"}}})107        ]108        deduped, _skipped = assignment.dedupe_latest(records)109 110        result = assignment.saved_responses_for_round(deduped, "adi", ["v1", "v2"])111 112        self.assertEqual(result, {}, "a completion from a prior round must not leak into this round's resume state")113 114    def test_empty_deduped_returns_empty_dict(self):115        self.assertEqual(assignment.saved_responses_for_round({}, "adi", ["v1"]), {})116 117 118class PickBalancedTests(unittest.TestCase):119    def test_prefers_lowest_coverage(self):120        coverage = {"v1": 4, "v2": 0, "v3": 2}121        picked = assignment.pick_balanced(list(coverage), coverage.get, 2, random.Random(0), target=5)122        self.assertEqual(set(picked), {"v2", "v3"})123 124    def test_excludes_at_or_over_target_when_enough_under_target_remain(self):125        coverage = {"v1": 5, "v2": 5, "v3": 1}126        picked = assignment.pick_balanced(list(coverage), coverage.get, 1, random.Random(0), target=5)127        self.assertEqual(picked, ["v3"])128 129    def test_falls_back_to_at_target_videos_when_not_enough_under_target(self):130        coverage = {"v1": 5, "v2": 5, "v3": 1}131        picked = assignment.pick_balanced(list(coverage), coverage.get, 3, random.Random(0), target=5)132        self.assertEqual(set(picked), {"v1", "v2", "v3"})133 134    def test_never_returns_more_than_requested_or_duplicates(self):135        coverage = {f"v{i}": i % 3 for i in range(30)}136        picked = assignment.pick_balanced(list(coverage), coverage.get, 7, random.Random(1), target=5)137        self.assertEqual(len(picked), 7)138        self.assertEqual(len(set(picked)), 7)139 140 141class CompletionCodeTests(unittest.TestCase):142    def test_format(self):143        code = assignment.generate_completion_code()144        self.assertTrue(code.startswith("T2AV-"))145        suffix = code[len("T2AV-"):]146        self.assertEqual(len(suffix), 12)147        self.assertTrue(all(c in assignment.COMPLETION_CODE_ALPHABET for c in suffix))148 149    def test_codes_are_not_trivially_repeated(self):150        codes = {assignment.generate_completion_code() for _ in range(200)}151        self.assertEqual(len(codes), 200, "200 draws from a 36^12 space should never collide")152 153 154class BuildRoundTests(unittest.TestCase):155    def test_always_excludes_already_completed_never_seeds_with_them(self):156        """Core policy change: a round is exactly N NEW videos, never padded157        with history, regardless of how much history exists."""158        catalog = [f"v{i}" for i in range(1, 201)]  # 200 videos, like production159        already_completed = set(catalog[:51])  # mirrors real production data (Adi: 51 completed)160 161        round_record = assignment.build_round(162            "adi", "Adi", catalog, already_completed, {}, {}, random.Random(0), round_number=1,163        )164 165        self.assertEqual(round_record["actual_size"], 20)166        self.assertEqual(len(round_record["video_ids"]), 20)167        self.assertEqual(168            already_completed.intersection(round_record["video_ids"]), set(),169            "none of the 51 previously completed videos may appear in the new round",170        )171        self.assertEqual(len(set(round_record["video_ids"])), 20, "no duplicates within the round")172        self.assertEqual(round_record["round_number"], 1)173        self.assertEqual(round_record["requested_size"], 20)174 175    def test_one_historical_completion_gets_twenty_new_not_nineteen(self):176        catalog = [f"v{i}" for i in range(1, 201)]177        already_completed = {"v1"}178 179        round_record = assignment.build_round(180            "youngsun", "youngsun_0715", catalog, already_completed, {}, {}, random.Random(0), round_number=1,181        )182 183        self.assertEqual(round_record["actual_size"], 20)184        self.assertNotIn("v1", round_record["video_ids"])185 186    def test_brand_new_annotator_gets_exactly_twenty(self):187        catalog = [f"v{i}" for i in range(1, 201)]188        round_record = assignment.build_round(189            "new", "New", catalog, set(), {}, {}, random.Random(0), round_number=1,190        )191        self.assertEqual(round_record["actual_size"], 20)192 193    def test_reduced_size_when_not_enough_eligible_remain(self):194        catalog = [f"v{i}" for i in range(1, 26)]  # only 25 videos total195        already_completed = set(catalog[:15])  # 15 done, only 10 eligible left196 197        round_record = assignment.build_round(198            "x", "X", catalog, already_completed, {}, {}, random.Random(0), round_number=1,199        )200 201        self.assertEqual(round_record["actual_size"], 10)202        self.assertEqual(round_record["requested_size"], 20)203        self.assertEqual(len(round_record["video_ids"]), 10)204        self.assertEqual(already_completed.intersection(round_record["video_ids"]), set())205 206    def test_second_round_excludes_first_rounds_completions_too(self):207        """After Adi finishes 51 (history) + 20 (round 1) = 71, round 2 must208        exclude all 71, selecting from the remaining 129."""209        catalog = [f"v{i}" for i in range(1, 201)]210        historical = set(catalog[:51])211 212        round1 = assignment.build_round(213            "adi", "Adi", catalog, historical, {}, {}, random.Random(0), round_number=1,214        )215        after_round1 = historical | set(round1["video_ids"])216        self.assertEqual(len(after_round1), 71)217 218        round2 = assignment.build_round(219            "adi", "Adi", catalog, after_round1, {}, {}, random.Random(1), round_number=2,220        )221 222        self.assertEqual(round2["actual_size"], 20)223        self.assertEqual(after_round1.intersection(round2["video_ids"]), set())224        remaining_pool_size = len(catalog) - len(after_round1)225        self.assertEqual(remaining_pool_size, 129)226        self.assertEqual(round2["round_number"], 2)227 228    def test_existing_completions_affect_new_round_priority(self):229        catalog = ["v1", "v2", "v3"]230        completed_by_video_map = {"v1": {"p1", "p2", "p3", "p4"}}  # coverage 4, still < target 5231 232        round_record = assignment.build_round(233            "new", "New", catalog, set(), completed_by_video_map, {}, random.Random(0),234            round_number=1, videos_per_annotator=2,235        )236 237        self.assertNotIn("v1", round_record["video_ids"], "v2/v3 have lower coverage and must be preferred")238        self.assertEqual(set(round_record["video_ids"]), {"v2", "v3"})239 240 241class ReservedByVideoTests(unittest.TestCase):242    def test_counts_open_uncompleted_unexpired_reservations(self):243        now = datetime(2026, 1, 10, tzinfo=timezone.utc)244        assignments = [245            {"annotator_id": "a", "video_ids": ["v1", "v2"], "created_at": "2026-01-09T00:00:00Z"},246            {"annotator_id": "b", "video_ids": ["v1"], "created_at": "2026-01-09T00:00:00Z"},247        ]248        completed_by_annotator = {"a": {"v2"}}249 250        reserved = assignment.reserved_by_video(assignments, completed_by_annotator, now, ttl_seconds=7 * 24 * 3600)251 252        self.assertEqual(reserved.get("v1"), 2)253        self.assertNotIn("v2", reserved)254 255    def test_excludes_expired_reservations_without_needing_to_delete_anything(self):256        now = datetime(2026, 1, 10, tzinfo=timezone.utc)257        assignments = [{"annotator_id": "a", "video_ids": ["v1"], "created_at": "2025-01-01T00:00:00Z"}]258 259        reserved = assignment.reserved_by_video(assignments, {}, now, ttl_seconds=24 * 3600)260 261        self.assertEqual(reserved.get("v1", 0), 0)262 263    def test_excludes_the_requesting_annotator_from_their_own_reservations(self):264        now = datetime(2026, 1, 10, tzinfo=timezone.utc)265        assignments = [{"annotator_id": "a", "video_ids": ["v1"], "created_at": "2026-01-09T00:00:00Z"}]266 267        reserved = assignment.reserved_by_video(assignments, {}, now, ttl_seconds=24 * 3600, exclude_annotator="a")268 269        self.assertEqual(reserved.get("v1", 0), 0)270 271 272class SequentialBalanceTests(unittest.TestCase):273    def test_sequential_rounds_stay_reasonably_balanced(self):274        catalog = [f"v{i}" for i in range(1, 81)]  # 80 videos275        reserved_by_video_map = {}276        rng = random.Random(42)277 278        for i in range(10):  # 10 annotators * 20 slots = 200 slots over 80 videos279            round_record = assignment.build_round(280                f"person{i}", f"person{i}", catalog, set(), {}, reserved_by_video_map, rng, round_number=1,281            )282            for video_id in round_record["video_ids"]:283                reserved_by_video_map[video_id] = reserved_by_video_map.get(video_id, 0) + 1284 285        counts = list(reserved_by_video_map.values())286        self.assertEqual(len(reserved_by_video_map), 80, "every video should have been touched")287        self.assertLessEqual(max(counts) - min(counts), 2, "balancing should keep coverage tight")288 289 290class FakeRepoStore:291    """Stands in for the annotations dataset repo's assignments/ and292    completions/ prefixes."""293 294    def __init__(self):295        self.files = {}296 297    def read(self, path):298        return self.files.get(path)299 300    def write_kwargs(self, **kwargs):301        path = kwargs["path_in_repo"]302        content = json.loads(kwargs["path_or_fileobj"].getvalue().decode("utf-8"))303        self.files[path] = content304 305    def list_active_round_records(self):306        records = []307        for path, pointer in self.files.items():308            if not path.endswith("/current.json"):309                continue310            annotator_id = path.split("/")[1]311            round_path = f"assignments/{annotator_id}/rounds/{pointer['assignment_id']}.json"312            completion_path = f"completions/{annotator_id}/{pointer['assignment_id']}.json"313            round_record = self.files.get(round_path)314            if round_record is None or completion_path in self.files:315                continue316            records.append(round_record)317        return records318 319 320class RoundLifecycleIntegrationTests(unittest.TestCase):321    """Exercises main.get_or_create_current_round() with every HF network322    call replaced by an in-memory fake, so this stays fully offline."""323 324    def setUp(self):325        import main as backend_main326 327        self.backend_main = backend_main328        self.store = FakeRepoStore()329        self.catalog = [f"v{i}" for i in range(1, 201)]  # 200 videos, like production330        self.completed_records = []  # raw annotation records fed to both list_annotation_records variants331 332        patches = [333            mock.patch.object(backend_main, "_read_json_from_repo", side_effect=self.store.read),334            # Both the cached and fresh variants must be mocked - the completion335            # check in get_or_create_current_round() deliberately bypasses the336            # cache (fresh=True) so a real annotator sees their code promptly;337            # leaving list_annotation_records() unmocked would silently fall338            # through to a real network call here.339            mock.patch.object(340                backend_main, "list_annotation_records_cached", side_effect=lambda: self.completed_records341            ),342            mock.patch.object(343                backend_main, "list_annotation_records", side_effect=lambda: self.completed_records344            ),345            mock.patch.object(346                backend_main, "list_active_round_records", side_effect=self.store.list_active_round_records347            ),348            mock.patch.object(backend_main, "ensure_dataset_exists", return_value=None),349            mock.patch.object(backend_main, "ensure_dataset_configured", return_value=None),350            mock.patch.object(backend_main, "_current_catalog_video_ids", return_value=self.catalog),351            mock.patch.object(backend_main, "upload_with_retry", side_effect=self.store.write_kwargs),352        ]353        for patcher in patches:354            patcher.start()355            self.addCleanup(patcher.stop)356 357    def _complete_round(self, annotator_raw, round_record):358        """Simulates the annotator finishing every video in a round by359        appending matching annotation records - exactly what save_annotation360        would have produced."""361        now = datetime.now(timezone.utc)362        for video_id in round_record["video_ids"]:363            self.completed_records.append(make_record(annotator_raw, video_id, iso(now)))364 365    def test_brand_new_annotator_gets_round_one_of_twenty(self):366        result = self.backend_main.get_or_create_current_round("Scratch Tester")367        self.assertEqual(result["status"], "in_progress")368        self.assertEqual(result["round_number"], 1)369        self.assertEqual(len(result["video_ids"]), 20)370        self.assertEqual(len(set(result["video_ids"])), 20)371 372    def test_refresh_while_incomplete_returns_identical_round(self):373        first = self.backend_main.get_or_create_current_round("Scratch Tester")374        second = self.backend_main.get_or_create_current_round("Scratch Tester")375        self.assertEqual(first["video_ids"], second["video_ids"])376        self.assertEqual(first["assignment_id"], second["assignment_id"])377 378    def test_no_completion_code_at_nineteen_of_twenty(self):379        first = self.backend_main.get_or_create_current_round("Scratch Tester")380        for video_id in first["video_ids"][:19]:381            self.completed_records.append(make_record("Scratch Tester", video_id, iso(datetime.now(timezone.utc))))382 383        result = self.backend_main.get_or_create_current_round("Scratch Tester")384 385        self.assertEqual(result["status"], "in_progress")386        self.assertNotIn("completion_code", result)387 388    def test_brand_new_round_has_no_saved_annotations(self):389        result = self.backend_main.get_or_create_current_round("Scratch Tester")390        self.assertEqual(result["saved_annotations"], {})391 392    def test_resume_returns_saved_responses_for_completed_videos_only(self):393        first = self.backend_main.get_or_create_current_round("Scratch Tester")394        completed_ids = first["video_ids"][:3]395        for video_id in completed_ids:396            self.completed_records.append(397                make_record(398                    "Scratch Tester",399                    video_id,400                    iso(datetime.now(timezone.utc)),401                    {"annotations": {"responses": {"video": {"tcRel": {"label": "PASS", "rationale": ""}}}}},402                )403            )404 405        result = self.backend_main.get_or_create_current_round("Scratch Tester")406 407        self.assertEqual(result["status"], "in_progress")408        self.assertEqual(set(result["saved_annotations"].keys()), set(completed_ids))409        for video_id in completed_ids:410            self.assertEqual(result["saved_annotations"][video_id]["video"]["tcRel"]["label"], "PASS")411 412    def test_completion_code_persisted_once_at_twenty_of_twenty(self):413        first = self.backend_main.get_or_create_current_round("Scratch Tester")414        self._complete_round("Scratch Tester", first)415 416        result_a = self.backend_main.get_or_create_current_round("Scratch Tester")417        result_b = self.backend_main.get_or_create_current_round("Scratch Tester")418 419        self.assertEqual(result_a["status"], "completed")420        self.assertTrue(result_a["completion_code"].startswith("T2AV-"))421        self.assertEqual(result_a["completion_code"], result_b["completion_code"], "code must be stable on refresh")422        completion_files = [p for p in self.store.files if p.startswith("completions/")]423        self.assertEqual(len(completion_files), 1, "completion record written exactly once")424 425    def test_within_grace_period_shows_same_completed_round_not_a_new_one(self):426        first = self.backend_main.get_or_create_current_round("Scratch Tester")427        self._complete_round("Scratch Tester", first)428        self.backend_main.get_or_create_current_round("Scratch Tester")  # triggers completion-record creation429 430        # Directly age the completion record to just inside the grace window.431        completion_path = f"completions/scratch-tester/{first['assignment_id']}.json"432        recent = datetime.now(timezone.utc) - timedelta(433            seconds=assignment.COMPLETION_GRACE_PERIOD_SECONDS - 30434        )435        self.store.files[completion_path]["completed_at"] = iso(recent)436 437        result = self.backend_main.get_or_create_current_round("Scratch Tester")438 439        self.assertEqual(result["status"], "completed")440        self.assertEqual(result["round_number"], 1)441 442    def test_past_grace_period_auto_creates_next_round_with_no_extra_call(self):443        first = self.backend_main.get_or_create_current_round("Scratch Tester")444        self._complete_round("Scratch Tester", first)445        self.backend_main.get_or_create_current_round("Scratch Tester")  # triggers completion-record creation446 447        # Directly age the completion record past the grace window - simulates "came back later".448        completion_path = f"completions/scratch-tester/{first['assignment_id']}.json"449        stale = datetime.now(timezone.utc) - timedelta(450            seconds=assignment.COMPLETION_GRACE_PERIOD_SECONDS + 60451        )452        self.store.files[completion_path]["completed_at"] = iso(stale)453 454        result = self.backend_main.get_or_create_current_round("Scratch Tester")455 456        self.assertEqual(result["status"], "in_progress")457        self.assertEqual(result["round_number"], 2)458        self.assertEqual(set(result["video_ids"]).intersection(first["video_ids"]), set())459 460    def test_never_assigns_the_same_video_twice_across_rounds(self):461        first = self.backend_main.get_or_create_current_round("Scratch Tester")462        self._complete_round("Scratch Tester", first)463        self.backend_main.get_or_create_current_round("Scratch Tester")464 465        completion_path = f"completions/scratch-tester/{first['assignment_id']}.json"466        stale = datetime.now(timezone.utc) - timedelta(seconds=assignment.COMPLETION_GRACE_PERIOD_SECONDS + 60)467        self.store.files[completion_path]["completed_at"] = iso(stale)468 469        second = self.backend_main.get_or_create_current_round("Scratch Tester")470 471        all_seen = first["video_ids"] + second["video_ids"]472        self.assertEqual(len(all_seen), len(set(all_seen)))473 474    def test_different_annotators_each_get_their_own_round(self):475        a = self.backend_main.get_or_create_current_round("Person A")476        b = self.backend_main.get_or_create_current_round("Person B")477        self.assertEqual(len(a["video_ids"]), 20)478        self.assertEqual(len(b["video_ids"]), 20)479        round_files = [p for p in self.store.files if "/rounds/" in p]480        self.assertEqual(len(round_files), 2)481 482    def test_simultaneous_requests_do_not_create_duplicate_rounds(self):483        """The lock fully serializes get_or_create_current_round, so we484        simulate "simultaneous" by calling it back-to-back and asserting only485        one round/pointer combination was ever written - the real concurrency486        guarantee is exercised live against a scratch HF repo separately."""487        for _ in range(5):488            self.backend_main.get_or_create_current_round("Scratch Tester")489        round_files = [p for p in self.store.files if "/rounds/" in p]490        self.assertEqual(len(round_files), 1)491 492    def test_partial_records_never_mistakenly_counted_as_completed(self):493        first = self.backend_main.get_or_create_current_round("Scratch Tester")494        # An autosave/reservation-like record missing created_at for every video.495        for video_id in first["video_ids"]:496            self.completed_records.append({"id": "x", "video_id": video_id, "user": "Scratch Tester"})497 498        result = self.backend_main.get_or_create_current_round("Scratch Tester")499 500        self.assertEqual(result["status"], "in_progress")501 502 503if __name__ == "__main__":504    unittest.main()505