tools/test_expand_dmhy_dag_annotations.py

"""Tests for full source-list expansion of DAG annotation patches."""

from __future__ import annotations

import json
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path


def write_jsonl(path: Path, rows: list[dict]) -> None:
    path.write_text(
        "\n".join(json.dumps(row, ensure_ascii=False) for row in rows) + "\n",
        encoding="utf-8",
    )


class ExpandDmhyDagAnnotationsTests(unittest.TestCase):
    def run_expand(
        self,
        tmp: Path,
        dag: dict,
        source_rows: list[dict],
        patch_rows: list[dict],
        *,
        limit: int | None = None,
    ) -> tuple[list[dict], dict, subprocess.CompletedProcess[str]]:
        dag_path = tmp / "dag.json"
        source_path = tmp / "dmhy_list.jsonl"
        patches_path = tmp / "patches.jsonl"
        output_path = tmp / "expanded.jsonl"
        manifest_path = tmp / "expanded.manifest.json"
        dag_path.write_text(json.dumps(dag, ensure_ascii=False), encoding="utf-8")
        write_jsonl(source_path, source_rows)
        write_jsonl(patches_path, patch_rows)

        command = [
            sys.executable,
            "-m",
            "tools.expand_dmhy_dag_annotations",
            "--dag",
            str(dag_path),
            "--source-list",
            str(source_path),
            "--patches",
            str(patches_path),
            "--output",
            str(output_path),
            "--manifest-output",
            str(manifest_path),
        ]
        if limit is not None:
            command.extend(["--limit", str(limit)])
        run = subprocess.run(command, check=False, capture_output=True, text=True)
        self.assertEqual(run.returncode, 0, run.stderr)
        rows = [
            json.loads(line)
            for line in output_path.read_text(encoding="utf-8").splitlines()
            if line.strip()
        ]
        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
        return rows, manifest, run

    def test_source_list_expands_multiple_values_for_same_terminal(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            rows, manifest, _run = self.run_expand(
                Path(tmpdir),
                {
                    "nodes": [],
                    "terminals": [
                        {
                            "terminal_id": "t0",
                            "node_id": 1,
                            "prefix": "[ANi] Full Show - ",
                            "value_examples": ["[ANi] Full Show - 01 [1080P].mkv"],
                        }
                    ],
                },
                [
                    {"value": "[ANi] Full Show - 01 [1080P].mkv"},
                    {"value": "[ANi] Full Show - 02 [1080P].mkv"},
                    {"value": "[ANi] Full Show - 03 [1080P].mkv"},
                ],
                [
                    {
                        "unit_id": "u0",
                        "terminal_ids": ["t0"],
                        "episode_title_suffixes": [],
                        "media_suffixes": ["[1080P]"],
                        "title_candidates": ["Full Show"],
                        "source": "responses-dag-v1",
                        "fallback": False,
                        "status": "ok",
                    }
                ],
            )

            self.assertEqual([row["filename"] for row in rows], [
                "[ANi] Full Show - 01 [1080P].mkv",
                "[ANi] Full Show - 02 [1080P].mkv",
                "[ANi] Full Show - 03 [1080P].mkv",
            ])
            self.assertTrue(all(row["terminal_id"] == "t0" for row in rows))
            self.assertEqual(manifest["input_rows"], 3)
            self.assertEqual(manifest["matched_rows"], 3)
            self.assertEqual(manifest["output_rows"], 3)
            self.assertEqual(manifest["unique_terminals_matched"], 1)

    def test_unmatched_values_are_counted(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            rows, manifest, _run = self.run_expand(
                Path(tmpdir),
                {
                    "nodes": [],
                    "terminals": [
                        {
                            "terminal_id": "t0",
                            "node_id": 1,
                            "prefix": "Known Show - ",
                            "value_examples": ["Known Show - 01 [1080P].mkv"],
                        }
                    ],
                },
                [
                    {"value": "Known Show - 01 [1080P].mkv"},
                    {"value": "Other Show - 01 [1080P].mkv"},
                    {"value": ""},
                ],
                [
                    {
                        "unit_id": "u0",
                        "terminal_ids": ["t0"],
                        "episode_title_suffixes": [],
                        "media_suffixes": ["[1080P]"],
                        "title_candidates": [],
                        "source": "responses-dag-v1",
                        "fallback": False,
                        "status": "ok",
                    }
                ],
            )

            self.assertEqual(len(rows), 1)
            self.assertEqual(manifest["input_rows"], 3)
            self.assertEqual(manifest["matched_rows"], 1)
            self.assertEqual(manifest["unmatched_rows"], 2)
            self.assertEqual(manifest["status_counts"]["unmatched"], 2)

    def test_patch_with_multiple_terminal_ids_maps_to_each_terminal(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            rows, manifest, _run = self.run_expand(
                Path(tmpdir),
                {
                    "nodes": [],
                    "terminals": [
                        {
                            "terminal_id": "t0",
                            "node_id": 1,
                            "prefix": "Shared A - ",
                            "value_examples": ["Shared A - 01 [1080P].mkv"],
                        },
                        {
                            "terminal_id": "t1",
                            "node_id": 2,
                            "prefix": "Shared B - ",
                            "value_examples": ["Shared B - 01 [1080P].mkv"],
                        },
                    ],
                },
                [
                    {"value": "Shared A - 01 [1080P].mkv"},
                    {"value": "Shared B - 02 [1080P].mkv"},
                ],
                [
                    {
                        "unit_id": "dag-node-1",
                        "terminal_ids": ["t0", "t1"],
                        "episode_title_suffixes": ["Finale"],
                        "media_suffixes": ["[1080P]"],
                        "title_candidates": ["Shared"],
                        "source": "responses-dag-v1",
                        "fallback": False,
                        "status": "ok",
                    }
                ],
            )

            self.assertEqual([row["terminal_id"] for row in rows], ["t0", "t1"])
            self.assertEqual({row["unit_id"] for row in rows}, {"dag-node-1"})
            self.assertTrue(all(row["annotations"]["terminal_ids"] == ["t0", "t1"] for row in rows))
            self.assertTrue(all(row["episode_title_suffixes"] == ["Finale"] for row in rows))
            self.assertEqual(manifest["patch_covered_rows"], 2)
            self.assertEqual(manifest["unique_terminals_matched"], 2)

    def test_limit_scans_only_requested_source_rows(self) -> None:
        with tempfile.TemporaryDirectory() as tmpdir:
            rows, manifest, _run = self.run_expand(
                Path(tmpdir),
                {
                    "nodes": [],
                    "terminals": [
                        {
                            "terminal_id": "t0",
                            "node_id": 1,
                            "prefix": "Limit Show - ",
                            "value_examples": ["Limit Show - 01 [1080P].mkv"],
                        }
                    ],
                },
                [
                    {"value": "Limit Show - 01 [1080P].mkv"},
                    {"value": "Limit Show - 02 [1080P].mkv"},
                ],
                [
                    {
                        "unit_id": "u0",
                        "terminal_ids": ["t0"],
                        "episode_title_suffixes": [],
                        "media_suffixes": ["[1080P]"],
                        "title_candidates": [],
                        "source": "responses-dag-v1",
                        "fallback": False,
                        "status": "ok",
                    }
                ],
                limit=1,
            )

            self.assertEqual(len(rows), 1)
            self.assertEqual(manifest["input_rows"], 1)


if __name__ == "__main__":
    unittest.main()