tools/test_schema_v2_synthetic_augment.py

"""Integration smoke test for the schema v2 synthetic augment Rust tool."""

from __future__ import annotations

import json
import os
import re
import subprocess
import tempfile
import unittest
from pathlib import Path
from typing import Any


REPO_ROOT = Path(__file__).resolve().parents[1]
MANIFEST_PATH = Path("tools/schema_v2_synthetic_augment/Cargo.toml")
RECIPES_PATH = Path("reports/dmhy_template_recipes.full_top5000.seed.jsonl")
LABEL_SCHEMA_PATH = Path("label_schema.json")
NUMERIC_TITLE_SEEDS_PATH = Path("data/synthetic_numeric_titles.txt")
PATH_PREFIX_SEEDS_PATH = Path("data/synthetic_path_prefixes.txt")

CATEGORY_FIELDS = (
    "source",
    "source_type",
    "augmentation",
    "augmentation_type",
    "augment_type",
    "kind",
    "category",
    "row_type",
    "generator",
    "variant",
)
TAG_PATH_SEGMENTS = {"Gekijouban", "2004", "TV"}


def load_jsonl(path: Path) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    with path.open("r", encoding="utf-8") as handle:
        for line_no, line in enumerate(handle, 1):
            line = line.strip()
            if not line:
                continue
            try:
                row = json.loads(line)
            except json.JSONDecodeError as exc:
                raise AssertionError(f"{path}:{line_no}: invalid JSON") from exc
            if not isinstance(row, dict):
                raise AssertionError(f"{path}:{line_no}: row must be a JSON object")
            rows.append(row)
    return rows


def load_text_lines(path: Path) -> list[str]:
    return [
        line.strip()
        for line in path.read_text(encoding="utf-8").splitlines()
        if line.strip()
    ]


def label_entity(label: str) -> str | None:
    if label == "O":
        return None
    return label.split("-", 1)[1]


def has_entity(row: dict[str, Any], entity: str) -> bool:
    return any(label_entity(label) == entity for label in row["labels"])


def has_path_entity(row: dict[str, Any]) -> bool:
    return any(
        (entity := label_entity(label)) is not None and entity.startswith("PATH_")
        for label in row["labels"]
    )


def category_text(row: dict[str, Any]) -> str:
    values = []
    for field in CATEGORY_FIELDS:
        value = row.get(field)
        if isinstance(value, str):
            values.append(value)
    return " ".join(values).lower().replace("-", "_")


def is_path_aug_row(row: dict[str, Any]) -> bool:
    text = category_text(row)
    return "path_aug" in text or "path_prefix" in text or has_path_entity(row)


def is_path_series_row(row: dict[str, Any]) -> bool:
    return "path_series" in category_text(row)


def is_path_movie_row(row: dict[str, Any]) -> bool:
    return "path_movie" in category_text(row)


def is_path_special_row(row: dict[str, Any]) -> bool:
    return "path_special" in category_text(row)


def is_path_confuser_row(row: dict[str, Any]) -> bool:
    return "path_confuser" in category_text(row)


def is_numeric_title_row(row: dict[str, Any], numeric_title_seeds: list[str]) -> bool:
    text = category_text(row)
    if "numeric_title" in text:
        return True
    if is_path_aug_row(row):
        return False
    filename = row.get("filename")
    return isinstance(filename, str) and any(seed in filename for seed in numeric_title_seeds)


def span_range(filename: str, text: str, *, start_at: int = 0, end_at: int | None = None) -> range:
    start = filename.find(text, start_at, end_at)
    if start < 0:
        raise AssertionError(f"span {text!r} not found in {filename!r}")
    return range(start, start + len(text))


def span_labels(row: dict[str, Any], text: str, *, start_at: int = 0, end_at: int | None = None) -> list[str]:
    filename = row["filename"]
    indexes = span_range(filename, text, start_at=start_at, end_at=end_at)
    return [row["labels"][index] for index in indexes]


def assert_span_not_entity(
    testcase: unittest.TestCase,
    row: dict[str, Any],
    text: str,
    entity: str,
    *,
    start_at: int = 0,
    end_at: int | None = None,
) -> None:
    labels = span_labels(row, text, start_at=start_at, end_at=end_at)
    testcase.assertFalse(
        any(label_entity(label) == entity for label in labels),
        msg=f"{text!r} in {row['filename']!r} was labeled as {entity}: {labels}",
    )


def iter_path_segments(filename: str) -> list[tuple[str, int, int]]:
    segments: list[tuple[str, int, int]] = []
    start = 0
    for match in re.finditer(r"[\\/]", filename):
        end = match.start()
        if end > start:
            segments.append((filename[start:end], start, end))
        start = match.end()
    if start < len(filename):
        segments.append((filename[start:], start, len(filename)))
    return segments


class SchemaV2SyntheticAugmentSmokeTests(unittest.TestCase):
    tmpdir: tempfile.TemporaryDirectory[str]
    output_path: Path
    manifest_output_path: Path
    rows: list[dict[str, Any]]
    manifest: dict[str, Any]
    allowed_labels: set[str]
    numeric_title_seeds: list[str]

    @classmethod
    def setUpClass(cls) -> None:
        cls.tmpdir = tempfile.TemporaryDirectory(prefix="anifilebert_schema_v2_aug_")
        tmp_path = Path(cls.tmpdir.name)
        cls.output_path = tmp_path / "aug.jsonl"
        cls.manifest_output_path = tmp_path / "aug.manifest.json"
        cls.numeric_title_seeds = load_text_lines(REPO_ROOT / NUMERIC_TITLE_SEEDS_PATH)
        cls.allowed_labels = set(
            json.loads((REPO_ROOT / LABEL_SCHEMA_PATH).read_text(encoding="utf-8"))[
                "labels"
            ]
        )

        command = [
            "cargo",
            "run",
            "--manifest-path",
            str(MANIFEST_PATH),
            "--bin",
            "schema_v2_synthetic_augment",
            "--",
            "--recipes",
            str(RECIPES_PATH),
            "--label-schema-file",
            str(LABEL_SCHEMA_PATH),
            "--numeric-title-seeds",
            str(NUMERIC_TITLE_SEEDS_PATH),
            "--path-prefix-seeds",
            str(PATH_PREFIX_SEEDS_PATH),
            "--limit-templates",
            "80",
            "--max-rows",
            "2000",
            "--output",
            str(cls.output_path),
            "--manifest-output",
            str(cls.manifest_output_path),
        ]
        env = os.environ.copy()
        env["CARGO_TARGET_DIR"] = str(tmp_path / "cargo-target")
        result = subprocess.run(
            command,
            cwd=REPO_ROOT,
            env=env,
            capture_output=True,
            text=True,
            encoding="utf-8",
            errors="replace",
        )
        if result.returncode != 0:
            raise RuntimeError(
                "\n".join(
                    [
                        f"cargo run failed with exit code {result.returncode}",
                        "command: " + " ".join(command),
                        "stdout:",
                        result.stdout.strip(),
                        "stderr:",
                        result.stderr.strip(),
                    ]
                )
            )

        cls.rows = load_jsonl(cls.output_path)
        cls.manifest = json.loads(
            cls.manifest_output_path.read_text(encoding="utf-8")
        )

    @classmethod
    def tearDownClass(cls) -> None:
        cls.tmpdir.cleanup()

    def test_jsonl_rows_are_char_tokenized_and_strict_bio(self) -> None:
        self.assertGreater(len(self.rows), 0, "generator produced no rows")
        for row_number, row in enumerate(self.rows, 1):
            with self.subTest(row=row_number):
                self.assertIsInstance(row.get("filename"), str)
                self.assertEqual(row["tokens"], list(row["filename"]))
                self.assertEqual(len(row["tokens"]), len(row["labels"]))

                previous_entity: str | None = None
                for index, label in enumerate(row["labels"]):
                    self.assertIn(label, self.allowed_labels)
                    if label == "O":
                        previous_entity = None
                        continue
                    prefix, entity = label.split("-", 1)
                    self.assertIn(prefix, {"B", "I"})
                    if prefix == "I":
                        self.assertEqual(
                            previous_entity,
                            entity,
                            msg=(
                                f"orphan I-tag at row {row_number}, index {index}: "
                                f"{label} in {row['filename']!r}"
                            ),
                        )
                    previous_entity = entity

    def test_manifest_counts_match_data_categories(self) -> None:
        required_keys = {"generated_rows", "numeric_title_rows", "path_rows"}
        self.assertFalse(required_keys - self.manifest.keys())

        numeric_title_rows = [
            row for row in self.rows if is_numeric_title_row(row, self.numeric_title_seeds)
        ]
        path_aug_rows = [row for row in self.rows if is_path_aug_row(row)]
        path_series_rows = [row for row in self.rows if is_path_series_row(row)]
        path_movie_rows = [row for row in self.rows if is_path_movie_row(row)]
        path_special_rows = [row for row in self.rows if is_path_special_row(row)]
        path_confuser_rows = [row for row in self.rows if is_path_confuser_row(row)]

        self.assertGreater(len(numeric_title_rows), 0)
        self.assertGreater(len(path_aug_rows), 0)
        self.assertGreater(len(path_series_rows), 0)
        self.assertGreater(len(path_movie_rows), 0)
        self.assertGreater(len(path_special_rows), 0)
        self.assertGreater(len(path_confuser_rows), 0)
        self.assertEqual(self.manifest["generated_rows"], len(self.rows))
        self.assertEqual(self.manifest["numeric_title_rows"], len(numeric_title_rows))
        self.assertEqual(self.manifest["path_rows"], len(path_aug_rows))
        self.assertEqual(self.manifest["path_series_rows"], len(path_series_rows))
        self.assertEqual(self.manifest["path_movie_rows"], len(path_movie_rows))
        self.assertEqual(self.manifest["path_special_rows"], len(path_special_rows))
        self.assertEqual(self.manifest["path_confuser_rows"], len(path_confuser_rows))
        self.assertEqual(
            self.manifest["path_rows"],
            self.manifest["path_series_rows"]
            + self.manifest["path_movie_rows"]
            + self.manifest["path_special_rows"]
            + self.manifest["path_confuser_rows"],
        )

    def test_numeric_title_digits_are_not_episode(self) -> None:
        for title, numeric_part in (("91 Days", "91"), ("7-nin", "7")):
            matches = [row for row in self.rows if title in row["filename"]]
            self.assertGreater(len(matches), 0, f"missing numeric title sample {title!r}")
            for row in matches:
                title_start = row["filename"].find(title)
                title_end = title_start + len(title)
                assert_span_not_entity(
                    self,
                    row,
                    numeric_part,
                    "EPISODE",
                    start_at=title_start,
                    end_at=title_end,
                )

    def test_special_codes_do_not_create_episode_labels(self) -> None:
        special_cases = ("[NCOP2]", "[NCED]", "[PV][01]")
        seen: set[str] = set()
        for row in self.rows:
            filename = row["filename"]
            for special in special_cases:
                if special not in filename:
                    continue
                seen.add(special)
                assert_span_not_entity(self, row, special, "EPISODE")
                self.assertIn("B-SPECIAL", span_labels(row, special))
        self.assertEqual(
            seen,
            set(special_cases),
            msg=f"missing special-code samples: {sorted(set(special_cases) - seen)}",
        )

    def test_path_aug_labels_path_title_path_season_and_tag_directories(self) -> None:
        path_rows = [row for row in self.rows if is_path_aug_row(row)]
        self.assertTrue(any(has_entity(row, "PATH_TITLE_LATIN") for row in path_rows))
        self.assertTrue(any(has_entity(row, "PATH_SEASON") for row in path_rows))
        self.assertTrue(any(is_path_confuser_row(row) for row in path_rows))

        seen_tag_segments: set[str] = set()
        for row in path_rows:
            for segment, start, end in iter_path_segments(row["filename"]):
                if segment not in TAG_PATH_SEGMENTS:
                    continue
                seen_tag_segments.add(segment)
                labels = row["labels"][start:end]
                self.assertEqual(
                    labels,
                    ["B-TAG", *(["I-TAG"] * (len(segment) - 1))],
                    msg=f"path segment {segment!r} should be TAG in {row['filename']!r}",
                )
                self.assertFalse(
                    any(label_entity(label) == "PATH_SEASON" for label in labels),
                    msg=f"path segment {segment!r} should not be PATH_SEASON",
                )

        self.assertEqual(
            seen_tag_segments,
            TAG_PATH_SEGMENTS,
            msg=f"missing path tag samples: {sorted(TAG_PATH_SEGMENTS - seen_tag_segments)}",
        )


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