"""Smoke tests for annotated DMHY graph dataset helpers.""" from __future__ import annotations import tempfile import json import subprocess import sys import unittest from pathlib import Path from tools.annotate_dmhy_prefix_graph import normalize_generated_tokens from tools.annotate_dmhy_prefix_tree_agent import ( Args as PrefixTreeAgentArgs, annotate_terminal_with_llm, annotate_selected, enrich_patch, request_body, validate_strict_annotation, ) from tools.convert_annotated_dmhy_dataset import ( iter_validated_jsonl, validate_record, ) from tools.convert_to_char_dataset import convert_record class AnnotatedDmhyWorkflowTests(unittest.TestCase): def prefix_tree_agent_args(self) -> PrefixTreeAgentArgs: return PrefixTreeAgentArgs( graph=Path("graph.json"), source_list=Path("source.jsonl"), output=Path("out.jsonl"), patch_output=Path("patch.jsonl"), units_output=Path("units.jsonl"), limit=None, min_weight=None, only_needs_review=False, llm=True, base_url="http://example.test/v1", api_key="test-key", model="gpt-5.4-mini", max_requests=None, http_timeout=1, preserve_i_labels=False, examples_only=False, workers=8, retries=2, resume=False, reasoning_effort="medium", ) def prefix_selected(self, count: int = 3) -> list[tuple[int, dict, dict]]: selected: list[tuple[int, dict, dict]] = [] for index in range(count): terminal_id = f"t{index}" terminal = { "terminal_id": terminal_id, "prefix": f"Show {index} - ", "suffix_examples": [f"0{index} - Title [1080P]"], "value_examples": [f"Show {index} - 0{index} - Title [1080P].mkv"], "weight": index + 1, } patch = enrich_patch( { "terminal_id": terminal_id, "needs_llm_review": True, "episode_title_suffixes": [], "media_suffixes": ["[1080P]"], "title_candidates": [], "llm_label": None, "notes": "heuristic", "source": "heuristic-v1", } ) selected.append((index, terminal, patch)) return selected def test_prefix_tree_agent_max_requests_does_not_fallback_unprocessed_terminals(self) -> None: args = self.prefix_tree_agent_args() args = PrefixTreeAgentArgs(**{**args.__dict__, "max_requests": 1, "workers": 1}) selected = self.prefix_selected(3) def fake_annotate(_terminal, patch, _args): merged = dict(patch) merged.update({"source": "responses-v1", "status": "ok", "fallback": False, "errors": []}) return merged, False import tools.annotate_dmhy_prefix_tree_agent as agent original_annotate = agent.annotate_terminal_with_llm try: agent.annotate_terminal_with_llm = fake_annotate completed, resume_skipped, llm_requests, llm_fallbacks, pending_unprocessed = annotate_selected(selected, args) finally: agent.annotate_terminal_with_llm = original_annotate self.assertEqual([patch["terminal_id"] for _ord, _idx, _term, patch in completed], ["t0"]) self.assertEqual(resume_skipped, 0) self.assertEqual(llm_requests, 1) self.assertEqual(llm_fallbacks, 0) self.assertEqual(pending_unprocessed, 2) self.assertFalse(any(patch["fallback"] for _ord, _idx, _term, patch in completed)) def test_prefix_tree_agent_resume_continues_next_batch(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: patch_path = Path(tmpdir) / "patches.jsonl" patch_path.write_text( json.dumps( { "terminal_id": "t0", "unit_id": "t0", "terminal_ids": ["t0"], "episode_title_suffixes": ["old"], "media_suffixes": ["[1080P]"], "title_candidates": ["old"], "llm_label": "old", "notes": "previous", "source": "responses-v1", "status": "ok", "fallback": False, "errors": [], }, ensure_ascii=False, separators=(",", ":"), ) + "\n", encoding="utf-8", ) args = self.prefix_tree_agent_args() args = PrefixTreeAgentArgs( **{**args.__dict__, "patch_output": patch_path, "resume": True, "max_requests": 1, "workers": 1} ) selected = self.prefix_selected(3) def fake_annotate(_terminal, patch, _args): merged = dict(patch) merged.update( { "episode_title_suffixes": [patch["terminal_id"]], "source": "responses-v1", "status": "ok", "fallback": False, "errors": [], } ) return merged, False import tools.annotate_dmhy_prefix_tree_agent as agent original_annotate = agent.annotate_terminal_with_llm try: agent.annotate_terminal_with_llm = fake_annotate completed, resume_skipped, llm_requests, _llm_fallbacks, pending_unprocessed = annotate_selected( selected, args ) finally: agent.annotate_terminal_with_llm = original_annotate patches = [patch for _ord, _idx, _term, patch in completed] self.assertEqual([patch["terminal_id"] for patch in patches], ["t0", "t1"]) self.assertEqual(patches[0]["notes"], "previous") self.assertEqual(patches[1]["episode_title_suffixes"], ["t1"]) self.assertEqual(resume_skipped, 1) self.assertEqual(llm_requests, 1) self.assertEqual(pending_unprocessed, 1) def test_prefix_tree_agent_completed_order_is_deterministic(self) -> None: args = self.prefix_tree_agent_args() args = PrefixTreeAgentArgs(**{**args.__dict__, "max_requests": 3, "workers": 3}) selected = self.prefix_selected(3) def fake_annotate(_terminal, patch, _args): merged = dict(patch) merged.update({"source": "responses-v1", "status": "ok", "fallback": False, "errors": []}) return merged, False import tools.annotate_dmhy_prefix_tree_agent as agent original_annotate = agent.annotate_terminal_with_llm try: agent.annotate_terminal_with_llm = fake_annotate completed, _resume_skipped, _llm_requests, _llm_fallbacks, _pending_unprocessed = annotate_selected( selected, args ) finally: agent.annotate_terminal_with_llm = original_annotate self.assertEqual([patch["terminal_id"] for _ord, _idx, _term, patch in completed], ["t0", "t1", "t2"]) def test_prefix_tree_agent_writes_llm_failures_without_completing_patches(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: failure_path = Path(tmpdir) / "failures.jsonl" args = self.prefix_tree_agent_args() args = PrefixTreeAgentArgs( **{**args.__dict__, "max_requests": 1, "workers": 1, "failure_output": failure_path} ) selected = self.prefix_selected(2) def fake_annotate(_terminal, patch, _args): failed = dict(patch) failed.update( { "status": "fallback", "fallback": True, "errors": ["Responses API HTTP 429: busy"], } ) return failed, True import tools.annotate_dmhy_prefix_tree_agent as agent original_annotate = agent.annotate_terminal_with_llm try: agent.annotate_terminal_with_llm = fake_annotate completed, _resume_skipped, llm_requests, llm_fallbacks, pending_unprocessed = annotate_selected( selected, args ) finally: agent.annotate_terminal_with_llm = original_annotate self.assertEqual(completed, []) self.assertEqual(llm_requests, 1) self.assertEqual(llm_fallbacks, 1) self.assertEqual(pending_unprocessed, 2) failures = [ json.loads(line) for line in failure_path.read_text(encoding="utf-8").splitlines() if line.strip() ] self.assertEqual(len(failures), 1) self.assertEqual(failures[0]["terminal_id"], "t0") self.assertEqual(failures[0]["status"], "pending") self.assertEqual(failures[0]["source"], "responses-v1") self.assertIn("Responses API HTTP 429", failures[0]["errors"][0]) def test_generated_tokens_split_punctuation_and_use_b_only_labels(self) -> None: tokens, labels = normalize_generated_tokens( ["[ANi]", " ", "Title-Name", "07"], ["B-GROUP", "O", "I-TITLE", "B-EPISODE"], ) self.assertEqual(tokens, ["[", "ANi", "]", " ", "Title", "-", "Name", "07"]) self.assertEqual( labels, ["O", "B-GROUP", "O", "O", "B-TITLE", "O", "B-TITLE", "B-EPISODE"], ) self.assertTrue(all(label == "O" or label.startswith("B-") for label in labels)) def test_preserve_i_labels_keeps_i_on_non_separator_pieces(self) -> None: tokens, labels = normalize_generated_tokens( ["Title-Name"], ["I-TITLE"], preserve_i_labels=True, ) self.assertEqual(tokens, ["Title", "-", "Name"]) self.assertEqual(labels, ["I-TITLE", "O", "I-TITLE"]) def test_prefix_tree_agent_strict_schema_payload(self) -> None: args = self.prefix_tree_agent_args() terminal = { "terminal_id": "t0", "prefix": "Show - 01", "suffix_examples": [" [1080P][WEB-DL]"], "value_examples": ["Show - 01 [1080P][WEB-DL].mkv"], } patch = { "terminal_id": "t0", "episode_title_suffixes": [], "media_suffixes": ["[1080P]", "[WEB-DL]"], "title_candidates": [], "notes": "heuristic", } body = request_body(args=args, terminal=terminal, patch=patch, previous_error=None, mode="reasoning_json_schema") self.assertEqual(body["model"], "gpt-5.4-mini") self.assertEqual(body["reasoning"], {"effort": "medium"}) schema_format = body["text"]["format"] self.assertEqual(schema_format["type"], "json_schema") self.assertTrue(schema_format["strict"]) self.assertIn("terminal_id", schema_format["schema"]["required"]) self.assertIn("episode_title_suffixes", schema_format["schema"]["required"]) def test_prefix_tree_agent_reasoning_effort_is_configurable(self) -> None: args = self.prefix_tree_agent_args() args = PrefixTreeAgentArgs(**{**args.__dict__, "reasoning_effort": "xhigh"}) body = request_body( args=args, terminal={"terminal_id": "t0"}, patch={"terminal_id": "t0", "notes": ""}, previous_error=None, mode="reasoning_json_schema", ) self.assertEqual(body["reasoning"], {"effort": "xhigh"}) def test_prefix_tree_agent_terminal_id_mismatch_validation(self) -> None: with self.assertRaisesRegex(ValueError, "terminal_id mismatch"): validate_strict_annotation( { "terminal_id": "wrong", "episode_title_suffixes": [], "media_suffixes": [], "title_candidates": [], "llm_label": "uncertain", "notes": "bad id", }, expected_terminal_id="expected", ) def test_prefix_tree_agent_retries_after_terminal_id_mismatch(self) -> None: args = self.prefix_tree_agent_args() terminal = { "terminal_id": "t0", "prefix": "Show - 01", "suffix_examples": [" - Finale [1080P]"], "value_examples": ["Show - 01 - Finale [1080P].mkv"], } patch = { "terminal_id": "t0", "needs_llm_review": True, "episode_title_suffixes": [], "media_suffixes": ["[1080P]"], "title_candidates": [], "llm_label": None, "notes": "heuristic", "source": "heuristic-v1", } calls: list[dict] = [] def fake_post(body, _args): calls.append(body) terminal_id = "wrong" if len(calls) == 1 else "t0" return { "output_text": json.dumps( { "terminal_id": terminal_id, "episode_title_suffixes": ["Finale"], "media_suffixes": ["[1080P]"], "title_candidates": ["Finale"], "llm_label": "mixed", "notes": "validated", } ) } import tools.annotate_dmhy_prefix_tree_agent as agent original_post = agent.post_responses try: agent.post_responses = fake_post merged, used_fallback = annotate_terminal_with_llm(terminal, patch, args) finally: agent.post_responses = original_post self.assertFalse(used_fallback) self.assertEqual(len(calls), 2) self.assertIn("Previous attempt failed validation", calls[1]["input"][0]["content"]) self.assertEqual(merged["terminal_id"], "t0") self.assertEqual(merged["source"], "responses-v1") self.assertEqual(merged["episode_title_suffixes"], ["Finale"]) def test_prefix_tree_agent_cli_writes_validator_compatible_units(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: tmp = Path(tmpdir) graph_path = tmp / "graph.json" dataset_path = tmp / "dmhy_weak.prefix_sample.jsonl" patch_path = tmp / "patches.jsonl" units_path = tmp / "units.jsonl" graph_path.write_text( json.dumps( { "terminals": [ { "terminal_id": "t0", "node_id": 42, "prefix": "[ANi] Prefix Show - ", "digit_skeleton": "[ANi] Prefix Show - ##", "weight": 7, "value_examples": [ "[ANi] Prefix Show - 01 [1080P][WEB-DL].mkv" ], "suffix_examples": ["01 [1080P][WEB-DL]"], } ] }, ensure_ascii=False, ), encoding="utf-8", ) annotate = subprocess.run( [ sys.executable, "-m", "tools.annotate_dmhy_prefix_tree_agent", "--graph", str(graph_path), "--output", str(dataset_path), "--patch-output", str(patch_path), "--units-output", str(units_path), "--examples-only", "--workers", "1", ], check=False, capture_output=True, text=True, ) self.assertEqual(annotate.returncode, 0, annotate.stderr) units = [ json.loads(line) for line in units_path.read_text(encoding="utf-8").splitlines() if line.strip() ] 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]["unit_id"], "t0") self.assertEqual(units[0]["source_kind"], "prefix_tree") self.assertEqual(units[0]["source_id"], 42) self.assertEqual(units[0]["terminal_ids"], ["t0"]) self.assertEqual(units[0]["weight"], 7) self.assertEqual(units[0]["context"]["prefixes"], ["[ANi] Prefix Show - "]) self.assertEqual(units[0]["context"]["digit_skeletons"], ["[ANi] Prefix Show - ##"]) self.assertEqual(units[0]["context"]["edge_labels"], []) self.assertEqual( units[0]["examples"]["values"], ["[ANi] Prefix Show - 01 [1080P][WEB-DL].mkv"], ) self.assertEqual(units[0]["examples"]["suffixes"], ["01 [1080P][WEB-DL]"]) self.assertEqual(units[0]["expected_output"], {"schema_version": "dmhy-annotation-v1"}) patches = [ json.loads(line) for line in patch_path.read_text(encoding="utf-8").splitlines() if line.strip() ] self.assertEqual(patches[0]["unit_id"], "t0") self.assertEqual(patches[0]["terminal_ids"], ["t0"]) self.assertEqual(patches[0]["status"], "ok") self.assertEqual(patches[0]["errors"], []) self.assertFalse(patches[0]["fallback"]) validate_units = subprocess.run( [ sys.executable, "-m", "tools.validate_dmhy_annotation_pipeline", "validate-units", str(units_path), ], check=False, capture_output=True, text=True, ) self.assertEqual(validate_units.returncode, 0, validate_units.stdout + validate_units.stderr) validate_patches = subprocess.run( [ sys.executable, "-m", "tools.validate_dmhy_annotation_pipeline", "validate-patches", str(patch_path), "--units", str(units_path), ], check=False, capture_output=True, text=True, ) self.assertEqual(validate_patches.returncode, 0, validate_patches.stdout + validate_patches.stderr) def test_validation_rejects_embedded_punctuation(self) -> None: record = { "filename": "Title-Name 07", "tokens": ["Title-Name", "07"], "labels": ["B-TITLE", "B-EPISODE"], } with self.assertRaisesRegex(ValueError, "contains punctuation"): validate_record(record, Path("sample.jsonl"), 1) def test_validation_rejects_embedded_symbol_separator(self) -> None: record = { "filename": "Title 1920×1080 07", "tokens": ["Title", "1920×1080", "07"], "labels": ["B-TITLE", "B-RESOLUTION", "B-EPISODE"], } with self.assertRaisesRegex(ValueError, "contains punctuation"): validate_record(record, Path("sample.jsonl"), 1) def test_b_only_input_converts_to_char_i_labels(self) -> None: record = { "filename": "Title-Name 07", "tokens": ["Title", "-", "Name", " ", "07"], "labels": ["B-TITLE", "O", "B-TITLE", "O", "B-EPISODE"], } validate_record(record, Path("sample.jsonl"), 1) converted = convert_record(record) self.assertIn("I-TITLE", converted["labels"]) self.assertEqual(converted["tokens"][:5], ["T", "i", "t", "l", "e"]) def test_iter_validated_jsonl_accepts_generated_shape(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: path = Path(tmpdir) / "records.jsonl" path.write_text( '{"filename":"A 01","tokens":["A"," ","01"],"labels":["B-TITLE","O","B-EPISODE"]}\n', encoding="utf-8", ) rows = list(iter_validated_jsonl(path)) self.assertEqual(len(rows), 1) self.assertEqual(rows[0]["filename"], "A 01") def test_cli_smoke_annotate_then_convert_with_temp_files(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: tmp = Path(tmpdir) graph_path = tmp / "graph.json" dataset_path = tmp / "dmhy_weak.generated.jsonl" char_path = tmp / "dmhy_weak.generated_char.jsonl" vocab_path = tmp / "vocab.generated.char.json" manifest_path = tmp / "manifest.json" graph_path.write_text( json.dumps( { "terminals": [ { "terminal_id": "t0", "weight": 1, "value_examples": [ "[ANi] Test Show - 01 [1080P][WEB-DL].mkv" ], "suffix_examples": [" [1080P][WEB-DL]"], } ] }, ensure_ascii=False, ), encoding="utf-8", ) annotate = subprocess.run( [ sys.executable, "-m", "tools.annotate_dmhy_prefix_graph", "--graph", str(graph_path), "--output", str(dataset_path), "--patch-output", "", "--examples-only", ], check=False, capture_output=True, text=True, ) self.assertEqual(annotate.returncode, 0, annotate.stderr) rows = [ json.loads(line) for line in dataset_path.read_text(encoding="utf-8").splitlines() if line.strip() ] self.assertEqual(len(rows), 1) self.assertIn("annotations", rows[0]) self.assertEqual(rows[0]["tokens"][0], "[") self.assertEqual(rows[0]["labels"][0], "O") convert = subprocess.run( [ sys.executable, "-m", "tools.convert_annotated_dmhy_dataset", "--input", str(dataset_path), "--output", str(char_path), "--vocab-output", str(vocab_path), "--manifest-output", str(manifest_path), "--progress", "0", ], check=False, capture_output=True, text=True, ) self.assertEqual(convert.returncode, 0, convert.stderr) self.assertTrue(char_path.exists()) self.assertTrue(vocab_path.exists()) self.assertTrue(manifest_path.exists()) def test_cli_source_list_mode_expands_beyond_value_examples(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: tmp = Path(tmpdir) graph_path = tmp / "graph.json" source_path = tmp / "dmhy_list.jsonl" dataset_path = tmp / "dmhy_weak.generated.jsonl" graph_path.write_text( json.dumps( { "terminals": [ { "terminal_id": "t0", "prefix": "[ANi] Full Show - ", "weight": 10, "value_examples": [ "[ANi] Full Show - 01 [1080P][WEB-DL].mkv" ], "suffix_examples": ["01 [1080P][WEB-DL]"], }, { "terminal_id": "t1", "prefix": "[ANi] Other Show - ", "weight": 10, "value_examples": [ "[ANi] Other Show - 01 [1080P][WEB-DL].mkv" ], "suffix_examples": ["01 [1080P][WEB-DL]"], }, ] }, ensure_ascii=False, ), encoding="utf-8", ) source_path.write_text( "\n".join( json.dumps({"value": value}, ensure_ascii=False) for value in [ "[ANi] Full Show - 01 [1080P][WEB-DL].mkv", "[ANi] Full Show - 02 [1080P][WEB-DL].mkv", "[ANi] Full Show - 03 [1080P][WEB-DL].mkv", "[ANi] Other Show - 01 [1080P][WEB-DL].mkv", ] ) + "\n", encoding="utf-8", ) annotate = subprocess.run( [ sys.executable, "-m", "tools.annotate_dmhy_prefix_graph", "--graph", str(graph_path), "--source-list", str(source_path), "--output", str(dataset_path), "--patch-output", "", "--limit", "1", ], check=False, capture_output=True, text=True, ) self.assertEqual(annotate.returncode, 0, annotate.stderr) rows = [ json.loads(line) for line in dataset_path.read_text(encoding="utf-8").splitlines() if line.strip() ] self.assertEqual(len(rows), 3) self.assertEqual([row["filename"] for row in rows], [ "[ANi] Full Show - 01 [1080P][WEB-DL].mkv", "[ANi] Full Show - 02 [1080P][WEB-DL].mkv", "[ANi] Full Show - 03 [1080P][WEB-DL].mkv", ]) self.assertTrue(all(row["terminal_id"] == "t0" for row in rows)) def test_cli_examples_only_uses_terminal_value_examples(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: tmp = Path(tmpdir) graph_path = tmp / "graph.json" source_path = tmp / "dmhy_list.jsonl" dataset_path = tmp / "dmhy_weak.generated.jsonl" graph_path.write_text( json.dumps( { "terminals": [ { "terminal_id": "t0", "prefix": "[ANi] Example Show - ", "weight": 10, "value_examples": [ "[ANi] Example Show - 01 [1080P][WEB-DL].mkv" ], "suffix_examples": ["01 [1080P][WEB-DL]"], } ] }, ensure_ascii=False, ), encoding="utf-8", ) source_path.write_text( "\n".join( json.dumps({"value": value}, ensure_ascii=False) for value in [ "[ANi] Example Show - 01 [1080P][WEB-DL].mkv", "[ANi] Example Show - 02 [1080P][WEB-DL].mkv", ] ) + "\n", encoding="utf-8", ) annotate = subprocess.run( [ sys.executable, "-m", "tools.annotate_dmhy_prefix_graph", "--graph", str(graph_path), "--source-list", str(source_path), "--output", str(dataset_path), "--patch-output", "", "--examples-only", ], check=False, capture_output=True, text=True, ) self.assertEqual(annotate.returncode, 0, annotate.stderr) rows = [ json.loads(line) for line in dataset_path.read_text(encoding="utf-8").splitlines() if line.strip() ] self.assertEqual(len(rows), 1) self.assertEqual(rows[0]["filename"], "[ANi] Example Show - 01 [1080P][WEB-DL].mkv") def test_cli_dag_annotation_units_include_shared_node_terminals(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: tmp = Path(tmpdir) dag_path = tmp / "dmhy_prefix_dag.json" output_path = tmp / "dmhy_prefix_dag.annotation_units.jsonl" dag_path.write_text( json.dumps( { "meta": {"version": "prefix-dag-v1"}, "root": 0, "nodes": [ { "id": 0, "terminal": False, "children": [ {"label": "A", "target": 1}, {"label": "B", "target": 2}, ], "incoming_count": 0, "reachable_terminals": 2, "reachable_weight": 20, }, { "id": 1, "terminal": False, "children": [{"label": " shared", "target": 3}], "incoming_count": 1, "reachable_terminals": 1, "reachable_weight": 10, }, { "id": 2, "terminal": False, "children": [{"label": " shared", "target": 3}], "incoming_count": 1, "reachable_terminals": 1, "reachable_weight": 10, }, { "id": 3, "terminal": False, "children": [ {"label": " 01", "target": 4}, {"label": " 02", "target": 5}, ], "incoming_count": 2, "reachable_terminals": 2, "reachable_weight": 20, }, { "id": 4, "terminal": True, "children": [], "incoming_count": 1, "reachable_terminals": 1, "reachable_weight": 10, }, { "id": 5, "terminal": True, "children": [], "incoming_count": 1, "reachable_terminals": 1, "reachable_weight": 10, }, ], "terminals": [ { "terminal_id": "t0", "node_id": 4, "prefix": "Show A shared 01", "digit_skeleton": "Show A shared ", "count": 10, "weight": 10, "suffix_examples": [" [1080P][WEB-DL]"], "value_examples": ["Show A shared 01 [1080P][WEB-DL].mkv"], "annotations": {}, }, { "terminal_id": "t1", "node_id": 5, "prefix": "Show B shared 02", "digit_skeleton": "Show B shared ", "count": 10, "weight": 10, "suffix_examples": [" [1080P][WEB-DL]"], "value_examples": ["Show B shared 02 [1080P][WEB-DL].mkv"], "annotations": {}, }, ], }, ensure_ascii=False, ), encoding="utf-8", ) annotate = subprocess.run( [ sys.executable, "-m", "tools.annotate_dmhy_prefix_dag", "--dag", str(dag_path), "--output", str(output_path), "--min-reachable-terminals", "2", "--min-incoming-count", "2", "--limit", "1", ], check=False, capture_output=True, text=True, ) self.assertEqual(annotate.returncode, 0, annotate.stderr) rows = [ json.loads(line) for line in output_path.read_text(encoding="utf-8").splitlines() if line.strip() ] self.assertEqual(len(rows), 1) self.assertEqual(rows[0]["unit_id"], "dag-node-3") self.assertEqual(rows[0]["kind"], "shared_suffix") self.assertEqual(rows[0]["terminal_ids"], ["t0", "t1"]) self.assertEqual( rows[0]["prefix_examples"], ["Show A shared 01", "Show B shared 02"], ) self.assertEqual(rows[0]["common_edge_labels"], [" 01", " 02"]) self.assertIn("annotations", rows[0]) def test_cli_dag_annotation_units_skip_terminal_sink_by_default(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: tmp = Path(tmpdir) dag_path = tmp / "dmhy_prefix_dag.json" output_path = tmp / "dmhy_prefix_dag.annotation_units.jsonl" dag_path.write_text( json.dumps( { "meta": {"version": "prefix-dag-v1"}, "root": 0, "nodes": [ { "id": 0, "terminal": False, "children": [{"label": "shared", "target": 2}], "incoming_count": 0, "reachable_terminals": 1, "reachable_weight": 12908, }, { "id": 2, "terminal": True, "children": [], "incoming_count": 12908, "reachable_terminals": 1, "reachable_weight": 12908, }, ], "terminals": [ { "terminal_id": "t-sink", "node_id": 2, "prefix": "Shared Sink", "digit_skeleton": "Shared Sink", "count": 12908, "weight": 12908, "suffix_examples": [], "value_examples": ["Shared Sink.mkv"], "annotations": {}, } ], }, ensure_ascii=False, ), encoding="utf-8", ) annotate = subprocess.run( [ sys.executable, "-m", "tools.annotate_dmhy_prefix_dag", "--dag", str(dag_path), "--output", str(output_path), "--min-reachable-terminals", "1", "--min-incoming-count", "2", ], check=False, capture_output=True, text=True, ) self.assertEqual(annotate.returncode, 0, annotate.stderr) rows = [ json.loads(line) for line in output_path.read_text(encoding="utf-8").splitlines() if line.strip() ] self.assertEqual(rows, []) def test_cli_dag_annotation_units_include_terminal_sink_when_requested(self) -> None: with tempfile.TemporaryDirectory() as tmpdir: tmp = Path(tmpdir) dag_path = tmp / "dmhy_prefix_dag.json" output_path = tmp / "dmhy_prefix_dag.annotation_units.jsonl" dag_path.write_text( json.dumps( { "meta": {"version": "prefix-dag-v1"}, "root": 0, "nodes": [ { "id": 0, "terminal": False, "children": [{"label": "shared", "target": 2}], "incoming_count": 0, "reachable_terminals": 1, "reachable_weight": 12908, }, { "id": 2, "terminal": True, "children": [], "incoming_count": 12908, "reachable_terminals": 1, "reachable_weight": 12908, }, ], "terminals": [ { "terminal_id": "t-sink", "node_id": 2, "prefix": "Shared Sink", "digit_skeleton": "Shared Sink", "count": 12908, "weight": 12908, "suffix_examples": [], "value_examples": ["Shared Sink.mkv"], "annotations": {}, } ], }, ensure_ascii=False, ), encoding="utf-8", ) annotate = subprocess.run( [ sys.executable, "-m", "tools.annotate_dmhy_prefix_dag", "--dag", str(dag_path), "--output", str(output_path), "--min-reachable-terminals", "1", "--min-incoming-count", "2", "--include-terminal-sink-units", ], check=False, capture_output=True, text=True, ) self.assertEqual(annotate.returncode, 0, annotate.stderr) rows = [ json.loads(line) for line in output_path.read_text(encoding="utf-8").splitlines() if line.strip() ] self.assertEqual(len(rows), 1) self.assertEqual(rows[0]["unit_id"], "dag-node-2") self.assertEqual(rows[0]["kind"], "shared_suffix") self.assertEqual(rows[0]["terminal_ids"], ["t-sink"]) self.assertEqual(rows[0]["common_edge_labels"], []) if __name__ == "__main__": unittest.main()