"""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)