158 lines
4.8 KiB
Python
158 lines
4.8 KiB
Python
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
|
|
|
|
|
|
def _make_srt(tmp_path, count=5) -> Path:
|
|
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:
|
|
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:
|
|
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_invoke_success(tmp_path, monkeypatch) -> None:
|
|
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:
|
|
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:
|
|
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:
|
|
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:
|
|
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"
|