"""Smoke tests for the independent DMHY DAG annotation runner.""" from __future__ import annotations import json import subprocess import sys import tempfile import unittest from pathlib import Path import tools.annotate_dmhy_dag_agent as agent def make_args(tmp: Path | None = None, **overrides) -> agent.Args: values = { "dag": Path("dag.json"), "output": Path("out.jsonl"), "patch_output": (tmp / "patch.jsonl") if tmp is not None else Path("patch.jsonl"), "units_output": Path("units.jsonl"), "limit": 10, "max_requests": 10, "units_per_request": 10, "max_context_units": 10, "workers": 8, "coverage_mode": agent.COVERAGE_SHARED_ONLY, "min_incoming_count": 2, "min_reachable_terminals": 2, "max_reachable_terminals": 500, "example_count": 2, "max_records": 20, "records_per_terminal": 1, "terminal_sink_threshold": 1000, "terminal_cluster_size": 1, "llm": True, "base_url": "http://example.test/v1", "api_key": "test-key", "model": "gpt-5.4-mini", "http_timeout": 1, "retries": 2, "preserve_i_labels": False, "resume": False, } values.update(overrides) return agent.Args(**values) def make_unit(unit_id: str, terminal_ids: list[str]) -> dict: terminals = [ { "_terminal_id": terminal_id, "_terminal_index": index, "value_examples": [f"Show {terminal_id} 01 [1080P].mkv"], "suffix_examples": ["01 [1080P]"], "weight": 1, } for index, terminal_id in enumerate(terminal_ids) ] return { "unit_id": unit_id, "node_id": 1, "kind": "shared_suffix", "incoming_count": 2, "reachable_terminals": len(terminal_ids), "reachable_weight": len(terminal_ids), "terminal_ids": terminal_ids, "prefix_examples": ["Show "], "digit_skeleton_examples": [], "value_examples": [terminal["value_examples"][0] for terminal in terminals], "suffix_examples": ["01 [1080P]"], "common_edge_labels": ["01"], "_terminals": terminals, } def response_for(units: list[dict], wrong_first_id: str | None = None) -> dict: annotations = [] for index, unit in enumerate(units): annotations.append( { "unit_id": wrong_first_id if index == 0 and wrong_first_id is not None else unit["unit_id"], "terminal_ids": unit["terminal_ids"], "episode_title_suffixes": [], "media_suffixes": ["[1080P]"], "title_candidates": [], "llm_label": "media", "notes": "ok", } ) return {"output_text": json.dumps({"annotations": annotations}, ensure_ascii=False)} class AnnotateDmhyDagAgentTests(unittest.TestCase): def test_batch_request_schema_contains_annotations_array(self) -> None: args = make_args() units = [make_unit("u1", ["t1"]), make_unit("u2", ["t2"])] body = agent.request_variants(units, args)[0] schema = body["text"]["format"]["schema"] self.assertEqual(body["text"]["format"]["name"], "dmhy_dag_batch_annotation") self.assertIn("annotations", schema["required"]) self.assertEqual(schema["properties"]["annotations"]["items"]["required"][0], "unit_id") payload = json.loads(body["input"]) self.assertEqual([unit["unit_id"] for unit in payload["units"]], ["u1", "u2"]) def test_batch_mismatch_retries_with_validation_error(self) -> None: args = make_args(retries=2) units = [make_unit("u1", ["t1"]), make_unit("u2", ["t2"])] calls: list[dict] = [] original_post = agent.http_post_json original_sleep = agent.time.sleep def fake_post(body, _args): calls.append(body) return response_for(units, wrong_first_id="wrong") if len(calls) == 1 else response_for(units) try: agent.http_post_json = fake_post agent.time.sleep = lambda _seconds: None patches = agent.annotate_batch(units, args, use_llm=True) finally: agent.http_post_json = original_post agent.time.sleep = original_sleep self.assertEqual(len(calls), 2) self.assertIn("Previous attempt failed validation", calls[1]["instructions"]) self.assertEqual([patch["unit_id"] for patch in patches], ["u1", "u2"]) self.assertTrue(all(patch["source"] == agent.LLM_SOURCE for patch in patches)) self.assertTrue(all(not patch["fallback"] for patch in patches)) def test_resume_skips_existing_completed_units_and_keeps_order(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: tmp = Path(tmpdir) args = make_args(tmp, resume=True, units_per_request=10, max_context_units=10, workers=1) units = [make_unit("u1", ["t1"]), make_unit("u2", ["t2"])] args.patch_output.write_text( json.dumps( { "unit_id": "u1", "terminal_ids": ["t1"], "episode_title_suffixes": ["Finale"], "media_suffixes": [], "title_candidates": [], "llm_label": "title", "notes": "existing", "source": agent.LLM_SOURCE, "fallback": False, "status": "ok", "attempts": 1, "errors": [], }, ensure_ascii=False, ) + "\n", encoding="utf-8", ) calls: list[list[str]] = [] original_post = agent.http_post_json def fake_post(body, _args): payload = json.loads(body["input"]) calls.append([unit["unit_id"] for unit in payload["units"]]) return response_for([units[1]]) try: agent.http_post_json = fake_post completed, skipped, llm_requested, pending_unprocessed = agent.annotate_units(units, args) finally: agent.http_post_json = original_post self.assertEqual(skipped, 1) self.assertEqual(llm_requested, 1) self.assertEqual(pending_unprocessed, 0) self.assertEqual(calls, [["u2"]]) patches = [patch for _unit, patch in completed] self.assertEqual([patch["unit_id"] for patch in patches], ["u1", "u2"]) self.assertEqual(patches[0]["episode_title_suffixes"], ["Finale"]) def test_max_requests_partial_does_not_fallback_unprocessed_units(self) -> None: args = make_args(max_requests=1, units_per_request=2, max_context_units=2, workers=1) units = [make_unit("u1", ["t1"]), make_unit("u2", ["t2"]), make_unit("u3", ["t3"])] calls: list[list[str]] = [] original_post = agent.http_post_json def fake_post(body, _args): payload = json.loads(body["input"]) batch_units = [unit for unit in units if unit["unit_id"] in {item["unit_id"] for item in payload["units"]}] calls.append([unit["unit_id"] for unit in payload["units"]]) return response_for(batch_units) try: agent.http_post_json = fake_post completed, skipped, llm_requested, pending_unprocessed = agent.annotate_units(units, args) finally: agent.http_post_json = original_post self.assertEqual(skipped, 0) self.assertEqual(llm_requested, 1) self.assertEqual(pending_unprocessed, 1) self.assertEqual(calls, [["u1", "u2"]]) patches = [patch for _unit, patch in completed] self.assertEqual([patch["unit_id"] for patch in patches], ["u1", "u2"]) self.assertTrue(all(not patch["fallback"] for patch in patches)) def test_resume_continues_with_next_unfinished_batch(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: tmp = Path(tmpdir) args = make_args(tmp, resume=True, max_requests=1, units_per_request=1, max_context_units=1, workers=1) units = [make_unit("u1", ["t1"]), make_unit("u2", ["t2"]), make_unit("u3", ["t3"])] args.patch_output.write_text( json.dumps( { "unit_id": "u1", "terminal_ids": ["t1"], "episode_title_suffixes": [], "media_suffixes": ["[1080P]"], "title_candidates": [], "llm_label": "media", "notes": "existing", "source": agent.LLM_SOURCE, "fallback": False, "status": "ok", "attempts": 1, "errors": [], }, ensure_ascii=False, ) + "\n", encoding="utf-8", ) calls: list[list[str]] = [] original_post = agent.http_post_json def fake_post(body, _args): payload = json.loads(body["input"]) calls.append([unit["unit_id"] for unit in payload["units"]]) return response_for([units[1]]) try: agent.http_post_json = fake_post completed, skipped, llm_requested, pending_unprocessed = agent.annotate_units(units, args) finally: agent.http_post_json = original_post self.assertEqual(skipped, 1) self.assertEqual(llm_requested, 1) self.assertEqual(pending_unprocessed, 1) self.assertEqual(calls, [["u2"]]) self.assertEqual([unit["unit_id"] for unit, _patch in completed], ["u1", "u2"]) def test_cli_heuristic_sample_skips_terminal_sink_and_writes_dataset_rows(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: tmp = Path(tmpdir) dag_path = tmp / "dag.json" output_path = tmp / "dmhy_weak.dag_sample.jsonl" patch_path = tmp / "patch.jsonl" units_path = tmp / "units.jsonl" dag_path.write_text( json.dumps( { "root": 0, "nodes": [ { "id": 0, "children": [{"label": "[A] Test Show - 0", "target": 1}], }, { "id": 1, "terminal": False, "children": [{"label": "1 [1080P][WEB-DL]", "target": 2}], "incoming_count": 2, "reachable_terminals": 2, "reachable_weight": 2, }, { "id": 2, "terminal": True, "children": [], "incoming_count": 10, "reachable_terminals": 100, }, ], "terminals": [ { "terminal_id": "t0", "node_id": 1, "prefix": "[A] Test Show - 0", "weight": 1, "suffix_examples": ["1 [1080P][WEB-DL]"], "value_examples": ["[A] Test Show - 01 [1080P][WEB-DL].mkv"], }, { "terminal_id": "t1", "node_id": 1, "prefix": "[A] Test Show - 0", "weight": 1, "suffix_examples": ["2 [1080P][WEB-DL]"], "value_examples": ["[A] Test Show - 02 [1080P][WEB-DL].mkv"], }, ], }, ensure_ascii=False, ), encoding="utf-8", ) run = subprocess.run( [ sys.executable, "-m", "tools.annotate_dmhy_dag_agent", "--dag", str(dag_path), "--output", str(output_path), "--patch-output", str(patch_path), "--units-output", str(units_path), "--limit", "1", "--max-requests", "0", "--workers", "2", ], check=False, capture_output=True, text=True, ) self.assertEqual(run.returncode, 0, run.stderr) patches = [json.loads(line) for line in patch_path.read_text(encoding="utf-8").splitlines()] self.assertEqual(len(patches), 1) self.assertEqual(patches[0]["unit_id"], "dag-node-1") self.assertEqual(patches[0]["terminal_ids"], ["t0", "t1"]) units = [json.loads(line) for line in units_path.read_text(encoding="utf-8").splitlines()] self.assertEqual(len(units), 1) self.assertEqual( set(units[0]), { "unit_id", "source_kind", "source_id", "terminal_ids", "weight", "context", "examples", "expected_output", }, ) self.assertEqual(units[0]["source_kind"], "prefix_dag") self.assertEqual(units[0]["source_id"], "1") self.assertEqual(units[0]["terminal_ids"], ["t0", "t1"]) self.assertEqual(units[0]["weight"], 2) self.assertEqual(units[0]["context"]["prefixes"], ["[A] Test Show - 0"]) self.assertEqual(units[0]["context"]["digit_skeletons"], []) self.assertEqual(units[0]["context"]["edge_labels"], ["1 [1080P][WEB-DL]"]) self.assertIn("incoming_count=2", units[0]["context"]["notes"]) self.assertEqual( units[0]["examples"]["values"], [ "[A] Test Show - 01 [1080P][WEB-DL].mkv", "[A] Test Show - 02 [1080P][WEB-DL].mkv", ], ) self.assertEqual(units[0]["examples"]["suffixes"], ["1 [1080P][WEB-DL]", "2 [1080P][WEB-DL]"]) self.assertEqual(units[0]["expected_output"], {"schema_version": "dmhy-annotation-v1"}) rows = [json.loads(line) for line in output_path.read_text(encoding="utf-8").splitlines()] self.assertEqual(len(rows), 2) self.assertEqual(rows[0]["tokens"][0], "[") self.assertTrue(all(label == "O" or label.startswith("B-") for row in rows for label in row["labels"])) def test_cli_all_terminals_adds_terminal_cluster_for_uncovered_terminal(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: tmp = Path(tmpdir) dag_path = tmp / "dag.json" output_path = tmp / "dmhy_weak.dag_all.jsonl" patch_path = tmp / "patch.jsonl" units_path = tmp / "units.jsonl" dag_path.write_text( json.dumps( { "root": 0, "nodes": [ { "id": 0, "children": [ {"label": "A shared", "target": 1}, {"label": "B solo", "target": 4}, ], }, { "id": 1, "children": [ {"label": " 01", "target": 2}, {"label": " 02", "target": 3}, ], "incoming_count": 2, "reachable_terminals": 2, "reachable_weight": 20, }, { "id": 2, "children": [], "incoming_count": 1, "reachable_terminals": 1, "reachable_weight": 10, }, { "id": 3, "children": [], "incoming_count": 1, "reachable_terminals": 1, "reachable_weight": 10, }, { "id": 4, "children": [], "incoming_count": 1, "reachable_terminals": 1, "reachable_weight": 5, }, ], "terminals": [ { "terminal_id": "t0", "node_id": 2, "prefix": "Shared Show 01", "weight": 10, "suffix_examples": [" [1080P]"], "value_examples": ["Shared Show 01 [1080P].mkv"], }, { "terminal_id": "t1", "node_id": 3, "prefix": "Shared Show 02", "weight": 10, "suffix_examples": [" [1080P]"], "value_examples": ["Shared Show 02 [1080P].mkv"], }, { "terminal_id": "t2", "node_id": 4, "prefix": "Solo Show 03", "weight": 5, "suffix_examples": [" [WEB-DL]"], "value_examples": ["Solo Show 03 [WEB-DL].mkv"], }, ], }, ensure_ascii=False, ), encoding="utf-8", ) run = subprocess.run( [ sys.executable, "-m", "tools.annotate_dmhy_dag_agent", "--dag", str(dag_path), "--output", str(output_path), "--patch-output", str(patch_path), "--units-output", str(units_path), "--coverage-mode", "all-terminals", "--min-incoming-count", "2", "--min-reachable-terminals", "2", "--limit", "100", "--max-requests", "0", ], check=False, capture_output=True, text=True, ) self.assertEqual(run.returncode, 0, run.stderr) summary = json.loads(run.stdout) self.assertEqual(summary["coverage_mode"], "all-terminals") self.assertEqual(summary["dag_terminals"], 3) self.assertEqual(summary["unique_terminal_coverage"], 3) patches = [json.loads(line) for line in patch_path.read_text(encoding="utf-8").splitlines()] self.assertEqual([patch["kind"] for patch in patches], ["shared_suffix", "terminal_cluster"]) self.assertEqual(patches[0]["terminal_ids"], ["t0", "t1"]) self.assertEqual(patches[1]["terminal_ids"], ["t2"]) units = [json.loads(line) for line in units_path.read_text(encoding="utf-8").splitlines()] self.assertEqual(len(units), 2) self.assertEqual(units[0]["unit_id"], "dag-node-1") self.assertEqual(units[0]["terminal_ids"], ["t0", "t1"]) self.assertEqual(units[1]["unit_id"], "terminal-cluster-1") self.assertEqual(units[1]["source_kind"], "prefix_dag") self.assertEqual(units[1]["source_id"], "terminal-cluster-1") self.assertEqual(units[1]["terminal_ids"], ["t2"]) self.assertIn("kind=terminal_cluster", units[1]["context"]["notes"]) self.assertEqual( {terminal_id for unit in units for terminal_id in unit["terminal_ids"]}, {"t0", "t1", "t2"}, ) rows = [json.loads(line) for line in output_path.read_text(encoding="utf-8").splitlines()] self.assertEqual({row["terminal_id"] for row in rows}, {"t0", "t1", "t2"}) self.assertIn("terminal-cluster-1", {row["unit_id"] for row in rows}) if __name__ == "__main__": unittest.main()