wov-node-llm/tests/test_llm_node.py

"""LLM 翻译节点测试。

覆盖真实 HTTP 服务调用、超时配置、SRT 回填、异常与入口点启动等路径。
"""

import json
import runpy
import threading
import urllib.error
import urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

from wov_node_llm.__main__ import invoke, translate_lines
from wov_sdk.models import InvokeRequest


class FakeUrlOpenResponse:
    """模拟 urllib 响应对象，提供固定 LLM 译文内容。"""

    def __init__(self, content: str) -> None:
        # 预编码为 Chat Completions 风格的 JSON 响应体。
        self._payload = json.dumps(
            {"choices": [{"message": {"content": content}}]}
        ).encode("utf-8")

    def read(self) -> bytes:
        return self._payload

    def __enter__(self):
        return self

    def __exit__(self, *args) -> bool:
        return False


def _make_srt(tmp_path, count=5) -> Path:
    """生成标准 SRT 测试文件，文本行为“原文字幕N”。"""
    lines = []
    for index in range(count):
        lines.extend(
            [
                str(index + 1),
                f"00:00:{index:02d},000 --> 00:00:{index + 1:02d},000",
                f"原文字幕{index + 1}",
                "",
            ]
        )
    path = tmp_path / "in.srt"
    path.write_text("\n".join(lines), encoding="utf-8")
    return path


def test_translate_lines_via_fake_api(monkeypatch) -> None:
    """验证通过真实 HTTP 服务器调用 LLM 接口并保持行顺序。"""
    class Handler(BaseHTTPRequestHandler):
        def do_POST(self) -> None:
            length = int(self.headers.get("Content-Length", "0"))
            self.rfile.read(length)
            body = json.dumps(
                {
                    "choices": [
                        {
                            "message": {
                                "content": "译文一\n译文二\n译文三\n译文四\n译文五"
                            }
                        }
                    ]
                }
            ).encode("utf-8")
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)

        def log_message(self, format, *args) -> None:
            return

    server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    try:
        monkeypatch.setenv(
            "LLM_API_BASE",
            f"http://127.0.0.1:{server.server_address[1]}/v1/chat/completions",
        )
        monkeypatch.setenv("LLM_API_KEY", "test-key")
        result = translate_lines(
            ["一", "二", "三", "四", "五"],
            {"target_language": "zh-CN"},
        )
        assert result == ["译文一", "译文二", "译文三", "译文四", "译文五"]
    finally:
        server.shutdown()
        server.server_close()
        thread.join(timeout=5)


def test_translate_lines_api_error(monkeypatch) -> None:
    """验证 LLM 接口不可用时抛出 URLError。"""
    def fail_open(request, timeout):
        raise urllib.error.URLError("api down")

    monkeypatch.setattr("wov_node_llm.__main__.urllib.request.urlopen", fail_open)
    try:
        translate_lines(["一"], {})
        raise AssertionError("expected failure")
    except urllib.error.URLError:
        pass


def test_translate_lines_default_timeout(monkeypatch) -> None:
    """验证未配置超时时使用默认 600 秒。"""
    captured = {}

    def fake_open(request, timeout):
        captured["timeout"] = timeout
        return FakeUrlOpenResponse("译文一")

    monkeypatch.setattr("wov_node_llm.__main__.urllib.request.urlopen", fake_open)
    monkeypatch.setenv("LLM_API_BASE", "http://fake/v1/chat/completions")

    result = translate_lines(["一"], {})

    assert result == ["译文一"]
    assert captured["timeout"] == 600


def test_translate_lines_env_timeout(monkeypatch) -> None:
    """验证 LLM_TIMEOUT_SECONDS 环境变量可覆盖超时。"""
    captured = {}

    def fake_open(request, timeout):
        captured["timeout"] = timeout
        return FakeUrlOpenResponse("译文一")

    monkeypatch.setattr("wov_node_llm.__main__.urllib.request.urlopen", fake_open)
    monkeypatch.setenv("LLM_API_BASE", "http://fake/v1/chat/completions")
    monkeypatch.setenv("LLM_TIMEOUT_SECONDS", "45")

    translate_lines(["一"], {})

    assert captured["timeout"] == 45


def test_invoke_success(tmp_path, monkeypatch) -> None:
    """验证成功调用会把译文回填到 SRT 并输出 cn.srt。"""
    source = _make_srt(tmp_path)

    def fake_translate(lines, params):
        return [f"译文{i + 1}" for i in range(len(lines))]

    monkeypatch.setattr("wov_node_llm.__main__.translate_lines", fake_translate)
    response = invoke(
        InvokeRequest(
            run_id="run_1",
            node_instance_id="ni_1",
            inputs={"srt_uri": str(source)},
            params={"target_language": "zh-CN"},
            output_dir=str(tmp_path / "out"),
        )
    )
    assert response.status == "completed"
    content = Path(response.outputs["cn_srt_uri"]).read_text(encoding="utf-8")
    assert "译文1" in content


def test_invoke_pads_short_translation(tmp_path, monkeypatch) -> None:
    """验证译文行数不足时用空行补齐，保持 SRT 结构完整。"""
    source = _make_srt(tmp_path, count=3)
    monkeypatch.setattr(
        "wov_node_llm.__main__.translate_lines",
        lambda lines, params: ["only one"],
    )
    response = invoke(
        InvokeRequest(
            run_id="run_2",
            node_instance_id="ni_2",
            inputs={"srt_uri": str(source)},
            output_dir=str(tmp_path / "out2"),
        )
    )
    assert response.status == "completed"


def test_invoke_missing_input(tmp_path) -> None:
    """验证缺少 srt_uri 时返回失败。"""
    response = invoke(
        InvokeRequest(
            run_id="run_3",
            node_instance_id="ni_3",
            inputs={},
            output_dir=str(tmp_path),
        )
    )
    assert response.status == "failed"


def test_invoke_missing_file(tmp_path) -> None:
    """验证 SRT 文件不存在时返回失败。"""
    response = invoke(
        InvokeRequest(
            run_id="run_4",
            node_instance_id="ni_4",
            inputs={"srt_uri": str(tmp_path / "missing.srt")},
            output_dir=str(tmp_path),
        )
    )
    assert response.status == "failed"


def test_entrypoint(monkeypatch) -> None:
    """验证 python -m wov_node_llm 会加载 llm-translate manifest。"""
    module_path = Path(__file__).resolve().parent.parent / "wov_node_llm" / "__main__.py"
    captured = {}

    def fake_run_node(manifest, handler) -> None:
        captured["id"] = manifest.id

    monkeypatch.setattr("wov_sdk.server.run_node", fake_run_node)
    runpy.run_path(str(module_path), run_name="__main__")
    assert captured["id"] == "llm-translate"