UmbraTrans/tools/translation/tests/test_translation.py

# Copyright 2026 Umbra Tools (umbratools.com)
# SPDX-License-Identifier: Apache-2.0

"""Unit tests for the translation cleanup and QA helpers.

The pipeline calls an LLM, so these tests deliberately cover only the
deterministic parts: prompt-leak removal, quote normalisation, English
residue detection, JSON extraction and the mechanical QA checks.
"""

from __future__ import annotations

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

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parent))

from book import SUBTITLES, _title_for  # noqa: E402
from languages import SCRIPT_LATIN, TargetLanguage, by_code, resolve  # noqa: E402
from llm import LLMError, parse_json  # noqa: E402
from pipeline import Runner, latin_residue, sanitize  # noqa: E402
from prompts import critic_prompt, draft_prompt, glossary_block, system_critic, system_drafter  # noqa: E402
from qa import check_chapter, forbidden_titles  # noqa: E402

CHINESE = resolve("zh-Hans")
SPANISH = resolve("es-ES")


class PromptInstructionTests(unittest.TestCase):
    """The drafter and the critic must both know about paragraph drift.

    Measured on the real Chinese edition: the drafter re-flows the passage
    across paragraph ids wherever Pass 0 split the original mid-sentence, and
    the alignment pass cannot see it - a one-to-one paragraph is a pure copy,
    so there is nothing to divide and nothing to disagree. The prevention is
    an instruction to the drafter; the detection is the alignment pass's
    ``--verify`` stage.
    """

    def _draft(self, target: TargetLanguage | None = None) -> str:
        return draft_prompt(
            chapter_title="Salvage",
            scene_index=0,
            chunk_index=0,
            chunks_in_scene=1,
            paragraphs=[{"id": "ch-6-p0000", "speaker": "narrator", "text": "The chair has"}],
            analysis={"summary": "摘要"},
            previous_summary="",
            glossary={},
            target=target or CHINESE,
        )

    def _critic(self, target: TargetLanguage | None = None) -> str:
        return critic_prompt(
            chapter_title="Salvage",
            paragraphs=[{"id": "ch-6-p0000", "speaker": "narrator", "text": "The chair has"}],
            draft={"ch-6-p0000": "中间的椅子旁"},
            glossary={},
            target=target or CHINESE,
        )

    def test_the_drafter_is_told_to_keep_each_id_to_its_own_english(self) -> None:
        prompt = self._draft()
        self.assertIn("每个 id 的译文只能覆盖该 id 自己的英文", prompt)

    def test_the_drafter_is_told_not_to_move_content_to_a_neighbour(self) -> None:
        prompt = self._draft()
        self.assertIn("不要把相邻段落的内容提前", prompt)
        self.assertIn("不要把本段的内容挪到下一个 id", prompt)

    def test_the_critic_is_asked_to_check_the_id_to_text_correspondence(self) -> None:
        prompt = self._critic()
        self.assertIn("misalignment", prompt)
        self.assertIn("段落 id 与译文是否对应", prompt)

    def test_a_misalignment_is_a_blocking_issue(self) -> None:
        prompt = self._critic()
        self.assertIn("段落错位", prompt)


class TargetLanguageTests(unittest.TestCase):
    """The pipeline must produce the language it was told to produce.

    Measured on oldbox: the production directory ``deal-hunter-es`` declared
    ``target_language: es-ES`` and held 3,038 paragraphs of Chinese, because
    the drafter, the critic and the glossary lookup were all hard-coded to
    Simplified Chinese. These tests pin the parameterisation that fixes it.
    """

    def test_a_language_is_found_by_tag_base_or_name(self) -> None:
        self.assertEqual(by_code("es-ES").base, "es")
        self.assertEqual(by_code("es").base, "es")
        self.assertEqual(by_code("es_ES").base, "es")
        self.assertEqual(by_code("Spanish").base, "es")
        self.assertEqual(by_code("简体中文").base, "zh")

    def test_an_unknown_language_falls_back_loudly(self) -> None:
        # The fallback prints a warning; the point is that it does not
        # silently pretend the language was recognised.
        self.assertEqual(resolve("xx-YY").base, "zh")

    def test_the_drafter_is_told_the_target_language(self) -> None:
        prompt = draft_prompt(
            chapter_title="Salvage",
            scene_index=0,
            chunk_index=0,
            chunks_in_scene=1,
            paragraphs=[{"id": "ch-6-p0000", "speaker": "narrator", "text": "The chair has"}],
            analysis={"summary": "摘要"},
            previous_summary="",
            glossary={},
            target=SPANISH,
        )
        self.assertIn("西班牙语", prompt)
        self.assertIn("español", prompt)
        self.assertNotIn("简体中文", prompt)

    def test_the_drafter_asks_for_the_target_language_key(self) -> None:
        prompt = draft_prompt(
            chapter_title="Salvage",
            scene_index=0,
            chunk_index=0,
            chunks_in_scene=1,
            paragraphs=[{"id": "ch-6-p0000", "speaker": "narrator", "text": "The chair has"}],
            analysis={"summary": "摘要"},
            previous_summary="",
            glossary={},
            target=SPANISH,
        )
        self.assertIn('"es":', prompt)
        self.assertNotIn('"zh":', prompt)

    def test_the_critic_is_told_the_target_language(self) -> None:
        prompt = critic_prompt(
            chapter_title="Salvage",
            paragraphs=[{"id": "ch-6-p0000", "speaker": "narrator", "text": "The chair has"}],
            draft={"ch-6-p0000": "la silla"},
            glossary={},
            target=SPANISH,
        )
        self.assertIn("西班牙语", prompt)
        self.assertNotIn("中文译文", prompt)

    def test_the_glossary_reads_the_target_language_key(self) -> None:
        glossary = {
            "title": {"en": "Deal Hunter"},
            "characters": {"Kainda": {"zh": "凯因达", "es": "Kainda"}},
            "terms": {"bounty": {"zh": "赏金", "es": "recompensa"}},
        }
        spanish = glossary_block(glossary, set(), "bounty", SPANISH)
        self.assertIn("recompensa", spanish)
        self.assertNotIn("赏金", spanish)

    def test_a_glossary_value_for_another_language_is_not_borrowed(self) -> None:
        """A Chinese-only glossary must not leak Chinese into a Spanish run."""
        glossary = {
            "characters": {"Kainda": {"zh": "凯因达"}},
            "terms": {},
        }
        spanish = glossary_block(glossary, set(), "", SPANISH)
        self.assertNotIn("凯因达", spanish)

    def test_the_chinese_rules_are_kept_for_chinese(self) -> None:
        """The existing production must not lose its hard rules."""
        glossary = {"characters": {}, "terms": {}, "places": {}}
        chinese = glossary_block(glossary, set(), "jump", CHINESE)
        self.assertIn("恒星系 system / star system = 星系", chinese)
        self.assertIn("不译成中文", chinese)

    def test_every_target_gets_the_books_own_rules(self) -> None:
        """The rules about this book are not Chinese's and not Spanish's.

        *Deal Hunter* is the ship's name, the AI's name and the book's title,
        so it stays in English in every edition. The Chinese table carries
        its own wording of this rule, so a language with no table must not
        lose it.
        """
        glossary = {"characters": {}, "terms": {}, "places": {}}
        for target in (CHINESE, SPANISH):
            block = glossary_block(glossary, set(), "", target)
            self.assertIn("Deal Hunter", block)
            self.assertIn("保留英文", block)
            self.assertIn("Hunter 作为 AI 的简称保留英文", block)
            self.assertIn("Umbra Audio 保留英文", block)

    def test_a_language_without_a_table_gets_no_other_languages_terms(self) -> None:
        """A Spanish run must not inherit the Chinese hard vocabulary."""
        glossary = {"characters": {}, "terms": {}, "places": {}}
        spanish = glossary_block(glossary, set(), "jump", SPANISH)
        self.assertNotIn("星系", spanish)
        self.assertNotIn("跃迁", spanish)

    def test_a_glossary_can_supply_its_own_rules(self) -> None:
        glossary = {
            "characters": {},
            "terms": {},
            "places": {},
            "rules": ["- Deal Hunter se conserva en inglés"],
        }
        spanish = glossary_block(glossary, set(), "", SPANISH)
        self.assertIn("Deal Hunter se conserva en inglés", spanish)

    def test_the_system_prompts_name_the_target_language(self) -> None:
        self.assertIn("简体中文", system_drafter(CHINESE))
        self.assertIn("西班牙语", system_drafter(SPANISH))
        self.assertIn("Spanish", system_critic(SPANISH))
        self.assertNotIn("中文译文", system_critic(SPANISH))

    def test_quotes_are_normalised_to_the_target_language(self) -> None:
        self.assertEqual(sanitize('"hola"', SPANISH), "“hola”")
        self.assertEqual(sanitize('"hallo"', resolve("de-DE")), "„hallo“")
        self.assertEqual(sanitize('"bonjour"', resolve("fr-FR")), "«bonjour»")

    def test_the_ellipsis_is_the_target_language_s(self) -> None:
        self.assertEqual(sanitize("...", CHINESE), "……")
        self.assertEqual(sanitize("...", SPANISH), "…")

    def test_the_residue_check_is_only_meaningful_for_a_non_latin_target(self) -> None:
        """A Spanish draft is full of Latin words; they are not residue."""
        self.assertIn("silla", latin_residue("la silla del capitán")[0])
        self.assertEqual(latin_residue("中间那把椅子")[0], [])
        # So the repair is skipped for a Latin-script target: the check would
        # otherwise flag the whole Spanish draft and "repair" it into
        # nonsense.
        self.assertIn(SCRIPT_LATIN, SPANISH.scripts)
        self.assertNotIn(SCRIPT_LATIN, CHINESE.scripts)


class RunDirectoryGuardTests(unittest.TestCase):
    """A run directory belongs to one language.

    Measured on oldbox: the production directory ``deal-hunter-es`` declared
    ``target_language: es-ES`` and held 3,038 paragraphs of Chinese, because
    the pipeline caches by file existence and the run directory is chosen by
    the caller. A Spanish run pointed at that directory would have skipped
    drafting entirely and published the Chinese under a Spanish name.

    The guard refuses that. It checks the manifest first (a claim) and then
    measures the text that is actually there (the evidence).
    """

    def _runner(self, run_dir: Path, target: str, **kwargs: object) -> Runner:
        book = run_dir / "book.json"
        glossary = run_dir / "glossary.json"
        book.write_text(
            json.dumps({"target_language": target, "chapters": []}),
            encoding="utf-8",
        )
        glossary.write_text(json.dumps({}), encoding="utf-8")
        return Runner(
            book_path=book,
            glossary_path=glossary,
            run_dir=run_dir / "runs",
            target_language=target,
            **kwargs,
        )

    def test_a_run_directory_records_its_language(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            self._runner(root, "es-ES")
            manifest = json.loads((root / "runs" / "run.json").read_text(encoding="utf-8"))
            self.assertEqual(manifest["target_language"], "es-ES")
            self.assertEqual(manifest["glossary_key"], "es")

    def test_a_spanish_run_refuses_a_chinese_run_directory(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            self._runner(root, "zh-Hans")
            with self.assertRaises(SystemExit) as caught:
                self._runner(root, "es-ES")
            self.assertIn("zh-Hans", str(caught.exception))
            self.assertIn("es-ES", str(caught.exception))

    def test_a_spanish_run_refuses_a_directory_holding_chinese_text(self) -> None:
        """No manifest - an older directory - so the text itself is measured."""
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            self._runner(root, "zh-Hans")
            chapter = root / "runs" / "ch-1"
            chapter.mkdir(parents=True, exist_ok=True)
            (chapter / "translation.json").write_text(
                json.dumps({"paragraphs": {"ch-1-p0000": "凯因达走到门前。"}}),
                encoding="utf-8",
            )
            (root / "runs" / "run.json").unlink()
            with self.assertRaises(SystemExit) as caught:
                self._runner(root, "es-ES")
            self.assertIn("not Spanish", str(caught.exception))

    def test_a_spanish_run_accepts_a_spanish_run_directory(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            self._runner(root, "es-ES")
            chapter = root / "runs" / "ch-1"
            chapter.mkdir(parents=True, exist_ok=True)
            (chapter / "translation.json").write_text(
                json.dumps({"paragraphs": {"ch-1-p0000": "Kainda se acercó a la puerta."}}),
                encoding="utf-8",
            )
            runner = self._runner(root, "es-ES")
            self.assertEqual(runner.target.code, "es-ES")

    def test_force_overrides_the_manifest(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            self._runner(root, "zh-Hans")
            runner = self._runner(root, "es-ES", force=True)
            self.assertEqual(runner.target.code, "es-ES")

    def test_a_draft_in_the_wrong_language_is_refused(self) -> None:
        """The drafter has been known to ignore the target language."""
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            runner = self._runner(root, "es-ES")
            with self.assertRaises(LLMError) as caught:
                runner._normalise_draft(
                    {"paragraphs": [{"id": "ch-1-p0000", "es": "凯因达走到门前。"}]},
                    ["ch-1-p0000"],
                )
            self.assertIn("wrong language", str(caught.exception))


class SanitizeTests(unittest.TestCase):
    def test_strips_dialogue_label(self) -> None:
        self.assertEqual(sanitize("（对话；凯因达）“我在游艇上。”"), "“我在游艇上。”")

    def test_strips_narration_label(self) -> None:
        self.assertEqual(sanitize("（旁白）凯因达登上了游艇。"), "凯因达登上了游艇。")

    def test_strips_paragraph_id(self) -> None:
        self.assertEqual(sanitize("[ch-6-p0000] “你好。”"), "“你好。”")

    def test_title_and_brand(self) -> None:
        self.assertEqual(sanitize("《交易猎人》"), "《Deal Hunter》")
        self.assertEqual(sanitize("由乌姆布拉音频系统制作"), "由Umbra Audio制作")

    def test_missing_hour(self) -> None:
        self.assertEqual(sanitize("她在小时前登船。"), "她一小时前登船。")

    def test_straight_quotes_become_chinese(self) -> None:
        self.assertEqual(sanitize('"你好"，她说。'), "“你好”，她说。")

    def test_unmatched_closing_quote_removed(self) -> None:
        self.assertEqual(sanitize("凯因达登上了游艇。”"), "凯因达登上了游艇。")

    def test_unmatched_opening_quote_closed(self) -> None:
        self.assertEqual(sanitize("“你好"), "“你好”")

    def test_ellipsis_and_dash(self) -> None:
        self.assertEqual(sanitize("等等... 然后--"), "等等…… 然后——")


class ResidueTests(unittest.TestCase):
    def test_lowercase_english_is_flagged(self) -> None:
        common, proper = latin_residue("她被 escorted 到了舱室。")
        self.assertEqual(common, ["escorted"])
        self.assertEqual(proper, [])

    def test_allowed_names_are_ignored(self) -> None:
        common, proper = latin_residue("Deal Hunter 和 Hunter 在 Umbra Audio。")
        self.assertEqual(common, [])
        self.assertEqual(proper, [])

    def test_unknown_capitalised_name_is_reported(self) -> None:
        common, proper = latin_residue("这是 Foo 的船。")
        self.assertEqual(common, [])
        self.assertEqual(proper, ["Foo"])


class JsonTests(unittest.TestCase):
    def test_plain_json(self) -> None:
        self.assertEqual(parse_json('{"a": 1}'), {"a": 1})

    def test_fenced_json(self) -> None:
        self.assertEqual(parse_json('```json\n{"a": 1}\n```'), {"a": 1})

    def test_json_with_prose(self) -> None:
        self.assertEqual(parse_json('结果如下：\n{"a": 1}\n完毕。'), {"a": 1})


class QaTests(unittest.TestCase):
    def test_missing_paragraph_is_block(self) -> None:
        chapter = {
            "id": "ch-x",
            "index": 1,
            "title": "Test",
            "paragraphs": [
                {"id": "ch-x-p0000", "type": "narration", "text": "One."},
                {"id": "ch-x-p0001", "type": "narration", "text": "Two."},
            ],
        }
        data = {"paragraphs": {"ch-x-p0000": "第一段。"}}
        issues, metrics = check_chapter(chapter, data)
        kinds = {i["kind"] for i in issues}
        self.assertIn("missing", kinds)
        self.assertEqual(metrics["translated"], 1)

    def test_metadata_leak_and_residue(self) -> None:
        chapter = {
            "id": "ch-x",
            "index": 1,
            "title": "Test",
            "paragraphs": [
                {"id": "ch-x-p0000", "type": "narration", "text": "One."},
                {"id": "ch-x-p0001", "type": "narration", "text": "The Albert system is far."},
            ],
        }
        data = {
            "paragraphs": {
                "ch-x-p0000": "（旁白）第一段 escorted。",
                "ch-x-p0001": "阿尔伯特系统很远。",
            }
        }
        issues, _ = check_chapter(chapter, data)
        kinds = {i["kind"] for i in issues}
        self.assertIn("metadata_leak", kinds)
        self.assertIn("english_residue", kinds)
        self.assertIn("system_mistranslated", kinds)

    def test_ship_systems_are_not_star_systems(self) -> None:
        chapter = {
            "id": "ch-x",
            "index": 1,
            "title": "Test",
            "paragraphs": [
                {"id": "ch-x-p0000", "type": "narration", "text": "Keep away from the ship systems."},
            ],
        }
        data = {"paragraphs": {"ch-x-p0000": "远离飞船系统。"}}
        issues, _ = check_chapter(chapter, data)
        self.assertNotIn("system_mistranslated", {i["kind"] for i in issues})


class ChapterTitleTests(unittest.TestCase):
    """The chapter title is the target language's, or nothing else's."""

    def test_a_language_with_no_table_falls_back_to_the_english_title(self) -> None:
        """Never to another language's title."""
        self.assertEqual(_title_for("Salvage", SPANISH), "Salvage")
        self.assertEqual(_title_for("Salvage", resolve("ja-JP")), "Salvage")

    def test_a_language_with_a_table_uses_it(self) -> None:
        self.assertEqual(_title_for("Salvage", CHINESE), "打捞")

    def test_a_chapter_missing_from_the_table_is_marked_in_the_target_language(self) -> None:
        """A missing heading is a visible gap, not an English title that looks done."""
        chinese = _title_for("Brand New Chapter", CHINESE)
        self.assertIn("待译", chinese)
        self.assertIn("Brand New Chapter", chinese)

    def test_the_subtitle_is_the_target_language_s(self) -> None:
        self.assertEqual(SUBTITLES["zh"], "交易猎人")
        self.assertEqual(SUBTITLES["es"], "Cazador de Ofertas")
        # A language with no subtitle gets none, not another language's.
        self.assertNotIn("en", SUBTITLES)


class ForbiddenTitleTests(unittest.TestCase):
    """The translated title must not appear in the text, in any language."""

    def test_the_chinese_form_is_forbidden(self) -> None:
        self.assertIn("《交易猎人》", forbidden_titles(CHINESE))

    def test_the_spanish_form_is_forbidden(self) -> None:
        self.assertIn("«Cazador de Ofertas»", forbidden_titles(SPANISH))

    def test_a_language_without_a_subtitle_forbids_nothing(self) -> None:
        """A language with no translated title has no forbidden form."""
        language = TargetLanguage(
            code="xx-XX",
            base="xx",
            name="Example",
            native="example",
            chinese_name="示例语",
            scripts=(SCRIPT_LATIN,),
        )
        self.assertEqual(forbidden_titles(language), [])


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