vrsub/tests/test_translation_line_alignment.py

"""翻译批处理行数对齐测试（先红后绿）。

背景：真实任务 run_51242078d76e（CJOD-255-长视频）产出的中文字幕存在
"内容-时间错位"——例如第 756 条「好像喜欢害羞的样子」被贴到 3805.34s
（该时间实际是日文「4つんばんですか(趴着吗)」的位置），而这条译文本应是
第 758 条「恥ずかしいのが好きみたいなので(喜欢害羞姿势)」的译文。

根因：nodes/llm.py 的 translate_lines 按 CHUNK_SIZE=20 分批把日文行发给
LLM，返回的译文行用 translated.extend() **无条件顺序拼接**，全批结束后只在
invoke 末尾做"多截断、少补空"。只要某批 LLM 返回行数 != 输入行数（实测大量
批次出现译文 21 行/原文 20 行），该批之后**所有字幕文本整体错位**，而时间戳
（从原文复制）保持不变 —— 造成"文本对错时间，程序从时间戳上看不出问题"。

修复（见 nodes/llm.py）：
1. 系统提示词新增"逐行独立翻译 + 碎片句按语境给含义 + 禁止合并/拆分"，
   从源头减少 LLM 重组断句导致的行数不一致；
2. 程序侧兜底 _repair_batch：返回行数 != 输入行数时，
   - 多行：末尾多余行合并到前一行（碎片本质同一句，时间轴保留）；
   - 少行：末尾补空串占位（宁缺勿错位，不挤占相邻字幕时间轴）。

本测试分两层：
1. _repair_batch / translate_lines 确定性单元测试（红 -> 绿）；
2. 真实数据 + 真实 LLM 集成测试（非 mock），验证产物与原文逐条对齐。
数据/Key 缺失时 skip。
"""

from __future__ import annotations

import json
import os
from pathlib import Path

import pytest

from tests.realdata_contract import parse_srt_entries

WORKSPACE = Path(__file__).resolve().parent.parent
TRANSCRIPT = Path(
    "/home/cat/Downloads/39.105.149.197/202609051737"
    "/run_51242078d76e/steps/asr/transcript.srt"
)

CHUNK_SIZE = 20


# ---------------------------------------------------------------------------
# 层一 helper：可注入的假 HTTP 客户端（与 nodes/llm.py 的 urllib 契约一致）
# ---------------------------------------------------------------------------


class _FakeUrlOpen:
    """模拟 urllib.request.urlopen：按调用次数依次返回预置的 LLM 输出。"""

    def __init__(self, contents: list[str]):
        self._contents = contents
        self._calls = 0

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        return False

    def read(self) -> bytes:
        content = self._contents[self._calls]
        self._calls += 1
        payload = {"choices": [{"message": {"content": content}}]}
        return json.dumps(payload).encode("utf-8")


def _patch_translate_llm(monkeypatch, batch_outputs: list[str]) -> None:
    """统一打桩：把 translate_lines 内 urlopen 换成 _FakeUrlOpen。"""
    import urllib.request

    # 直接替换 urllib.request.urlopen（nodes/llm.py 也是经它调用）。

    # 直接替换 urllib.request.urlopen（nodes/llm.py 也是经它调用）。
    fake = _FakeUrlOpen(batch_outputs)
    monkeypatch.setattr(urllib.request, "urlopen", lambda req, timeout=None: fake)
    monkeypatch.setenv("LLM_API_KEY", "test-key")
    monkeypatch.setenv("LLM_API_BASE", "http://fake/v1/chat/completions")


# ---------------------------------------------------------------------------
# 层一：_repair_batch 确定性单元测试
# ---------------------------------------------------------------------------


def test_repair_batch_extra_lines_merged() -> None:
    """多行：LLM 返回 21 行但输入 20 行，末尾多余行应合并到前一行。"""
    from nodes.llm import _repair_batch

    out = _repair_batch([f"译{i}" for i in range(21)], 20)
    assert len(out) == 20
    # 最后一行 = 原第 19（索引19）+第 20（索引20）行的合并。
    assert out[19] == "译19 译20"


def test_repair_batch_fewer_lines_padded() -> None:
    """少行：LLM 返回 19 行但输入 20 行，末尾补空串占位不挤占时间轴。"""
    from nodes.llm import _repair_batch

    out = _repair_batch([f"译{i}" for i in range(19)], 20)
    assert len(out) == 20
    assert out[19] == ""


def test_repair_batch_exact_unchanged() -> None:
    """正好对齐：原样返回。"""
    from nodes.llm import _repair_batch

    out = _repair_batch([f"译{i}" for i in range(20)], 20)
    assert len(out) == 20
    assert out == [f"译{i}" for i in range(20)]


# ---------------------------------------------------------------------------
# 层一：translate_lines 整批校验（多行/少行场景经修复后必须对齐）
# ---------------------------------------------------------------------------


@pytest.mark.integration
def test_translate_lines_aligns_extra_line(monkeypatch) -> None:
    """输入 40 行（两批 20），首批 LLM 返回 21 行：修复后必须对齐为 40 行。"""
    from nodes import llm as llm_node

    src_lines = [f"原文{i}" for i in range(40)]
    batch1_wrong = "\n".join([f"译{i}" for i in range(21)])  # 21 行错位源
    batch2_ok = "\n".join([f"译{i}" for i in range(20, 40)])
    _patch_translate_llm(monkeypatch, [batch1_wrong, batch2_ok])

    result = llm_node.translate_lines(src_lines, {})
    assert len(result) == len(src_lines), (
        f"translate_lines 未把多行合并对齐：输入 {len(src_lines)} 行，返回 {len(result)} 行"
    )


@pytest.mark.integration
def test_translate_lines_aligns_missing_line(monkeypatch) -> None:
    """第二批 LLM 少行时触发重试：重试返回正确 20 行后必须仍为 40 行。"""
    from nodes import llm as llm_node

    src_lines = [f"原文{i}" for i in range(40)]
    batch1_ok = "\n".join([f"译{i}" for i in range(20)])
    # 第二批第一次返回 19 行（少行）-> 触发重试；第二次返回正确 20 行。
    batch2_short = "\n".join([f"译{i}" for i in range(20, 39)])  # 19 行
    batch2_retry = "\n".join([f"译{i}" for i in range(20, 40)])  # 20 行
    _patch_translate_llm(monkeypatch, [batch1_ok, batch2_short, batch2_retry])

    result = llm_node.translate_lines(src_lines, {})
    assert len(result) == len(src_lines), (
        f"translate_lines 未把少行补齐：输入 {len(src_lines)} 行，返回 {len(result)} 行"
    )


@pytest.mark.integration
def test_translate_lines_pads_after_retries_exhausted(monkeypatch) -> None:
    """少行且重试耗尽：必须补空串占位，仍保持与输入等长（宁缺勿错位）。"""
    from nodes import llm as llm_node

    src_lines = [f"原文{i}" for i in range(40)]
    batch1_ok = "\n".join([f"译{i}" for i in range(20)])
    batch2_short = "\n".join([f"译{i}" for i in range(20, 39)])
    from nodes.llm import MAX_BATCH_RETRIES

    # 首次调用 + 重试重发，共 MAX_BATCH_RETRIES 次对 batch2 的调用都返回 19 行。
    responses = [batch1_ok] + [batch2_short] * MAX_BATCH_RETRIES
    _patch_translate_llm(monkeypatch, responses)

    result = llm_node.translate_lines(src_lines, {})
    assert len(result) == len(src_lines), (
        f"重试耗尽后未能补空串：输入 {len(src_lines)} 行，返回 {len(result)} 行"
    )


# ---------------------------------------------------------------------------
# 层二：真实数据 + 真实 LLM 集成测试（非 mock）
# ---------------------------------------------------------------------------


def _llm_credentials_ok() -> bool:
    """是否具备真实 LLM 调用条件（加载 .env 后 Key 非空）。"""
    try:
        from dotenv import load_dotenv

        load_dotenv(WORKSPACE / ".env")
    except Exception:
        pass
    return bool(os.getenv("LLM_API_KEY"))


@pytest.mark.integration
def test_pipeline_zh_cn_timetext_alignment(tmp_path) -> None:
    """真实数据 + 真实 LLM：完整翻译流水线后，译文必须与原文时间逐条对齐。

    方法：把真实日文 transcript.srt 喂给 llm.invoke（真实 LLM API），产出
    cn.srt；逐条比较 cn.srt 与原文的 (start, 行序) 严格一致。
    """
    if not _llm_credentials_ok():
        pytest.skip("未配置 LLM_API_KEY，跳过真实 LLM 集成测试")
    if not TRANSCRIPT.is_file():
        pytest.skip("缺少真实 transcript.srt，跳过集成测试")

    from wov_sdk.models import InvokeRequest
    from nodes import llm as llm_node

    out_dir = tmp_path / "out"
    response = llm_node.invoke(
        InvokeRequest(
            run_id="align_llm_test",
            node_instance_id="",
            inputs={"srt_uri": str(TRANSCRIPT)},
            params={"target_language": "zh-CN"},
            output_dir=str(out_dir),
        )
    )
    assert response.status == "completed", response.error

    zh_path = Path(response.outputs["cn_srt_uri"])
    zh_entries = parse_srt_entries(zh_path.read_text(encoding="utf-8"))
    src_entries = parse_srt_entries(TRANSCRIPT.read_text(encoding="utf-8"))
    assert len(zh_entries) == len(src_entries), (
        f"译文条数 {len(zh_entries)} != 原文 {len(src_entries)}：批内行数不一致导致错位。"
    )

    for i, (ze, se) in enumerate(zip(zh_entries, src_entries)):
        if abs(ze["start"] - se["start"]) > 0.01:
            raise AssertionError(
                f"第 {i} 条译文时间 {ze['start']:.2f} != 原文 {se['start']:.2f}："
                f"译文文本已整体错位（原文 '{se['text'][:15]}'）"
            )