按"测试规则"重写 tests/:一个模块一个目录、用例按数据→过程→验证三段书写、 不保留全局 conftest.py、测试过程只调用真实生产代码。 结构(73 个文件、30 个模块目录、477 用例): - tests/nodes/ 15 个模块目录(srt/whisper/ass/ffmpeg/frame_extract/vlm/ subtitle_ocr/llm/llm_filter/subtitle_cleanup/subtitle_correction/ proper_nouns/adaptive_pool/vad_profiler/echo); - tests/app/ 11 个模块目录(db/scheduler/batch/maintenance/registry/seed/ storage/config/logging/main/routers 三组 API); - tests/sdk/test_models、tests/web/test_crop、tests/shared(公共设施)。 测试数据随模块目录入库(tests/**/data/),删除根级 testdata/;.gitignore 的 data/ 改为 /data/,否则会连带忽略 tests/**/data/ 导致测试数据无法入库。 顺带发现并修复三个真实缺陷: - nodes/srt.py:相邻条目缺少空行时把下一条时间轴吞进正文(静默错位), 改为正文行遇时间戳行即报错; - src/wov_app/scheduler.py:_file_size 只捕获 OSError,含 \x00 的产物 URI 抛 ValueError 导致任务误判失败,改为同时捕获; - nodes/subtitle_correction.py:生产代码依赖测试包解析 SRT, 改用生产模块 nodes/srt.py。 真实模型/服务集成测试按外部状态跳过:新增 tests/shared/gpu_memory.py (运行时探测显存、CUDA OOM 转跳过)与 tests/shared/llm_service.py (无 Key / 余额 / 限流转跳过)。全量 477 passed。
411 lines
14 KiB
Python
411 lines
14 KiB
Python
"""nodes/vlm.py 的模块级测试(数据 → 测试过程 → 验证结果)。
|
||
|
||
被测模块:`nodes/vlm.py`(单帧 OCR:调 Ollama /api/chat 流式识别),可独立
|
||
调用。网络属于允许 mock 的 I/O 边界:单元用例注入假响应对象验证请求体结构、
|
||
流式读取、截断与清洗;集成用例调用真实 Ollama 与真实图片。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import socket
|
||
import urllib.error
|
||
import urllib.request
|
||
from pathlib import Path
|
||
|
||
import pytest
|
||
|
||
from nodes.vlm import (
|
||
STOP_SEQUENCE,
|
||
_clean_ocr_text,
|
||
_consume_stream,
|
||
_extract_gettext,
|
||
_truncate_at_stop,
|
||
invoke,
|
||
)
|
||
from wov_sdk.models import InvokeRequest
|
||
|
||
# 模块专用真实素材:有文字帧 / 无文字帧 / 真实视频字幕截图。
|
||
DATA_DIR = Path(__file__).resolve().parent / "data"
|
||
TEXT_IMAGE = DATA_DIR / "ocr_text.png"
|
||
NOTEXT_IMAGE = DATA_DIR / "ocr_notext.png"
|
||
REAL_SUBTITLE_IMAGE = DATA_DIR / "test_real_hav_sub.png"
|
||
|
||
# 集成测试期望识别出的真实字幕文本。
|
||
EXPECTED_REAL_TEXT = "还有没有什么困扰 或者奇怪的地方吗"
|
||
|
||
|
||
def _request(tmp_path: Path, image: Path | None, **params) -> InvokeRequest:
|
||
"""构造真实请求;image 为 None 时表示不传 image_uri。"""
|
||
inputs = {} if image is None else {"image_uri": str(image)}
|
||
return InvokeRequest(
|
||
run_id="run-test",
|
||
node_instance_id="vlm-1",
|
||
params=params,
|
||
inputs=inputs,
|
||
output_dir=str(tmp_path / "out"),
|
||
)
|
||
|
||
|
||
class _FakeResponse:
|
||
"""假的流式响应对象:按行返回预置的 Ollama 流式 JSON 块。"""
|
||
|
||
def __init__(self, lines: list[bytes]) -> None:
|
||
self._lines = list(lines)
|
||
|
||
def readline(self) -> bytes:
|
||
"""返回下一行;耗尽后返回空字节(模拟流结束)。"""
|
||
return self._lines.pop(0) if self._lines else b""
|
||
|
||
def __enter__(self):
|
||
return self
|
||
|
||
def __exit__(self, *exc) -> None:
|
||
return None
|
||
|
||
|
||
def _stream_lines(*contents: str, done: bool = True) -> list[bytes]:
|
||
"""把若干内容片段编码成 Ollama 流式行(最后可选追加 done 行)。"""
|
||
lines = [
|
||
json.dumps({"message": {"content": c}}).encode("utf-8") for c in contents
|
||
]
|
||
if done:
|
||
lines.append(json.dumps({"done": True}).encode("utf-8"))
|
||
return lines
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 文本清洗与标签提取
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_clean_ocr_text_strips_fences_and_blank_lines() -> None:
|
||
"""清洗去掉 markdown 围栏行与空行,只保留识别文字。"""
|
||
# 数据:glm-ocr 典型输出(识别文本后追加大量围栏)。
|
||
raw = "```markdown\n字幕第一行\n\n字幕第二行\n```\n"
|
||
|
||
# 测试过程
|
||
cleaned = _clean_ocr_text(raw)
|
||
|
||
# 验证结果
|
||
assert cleaned == "字幕第一行\n字幕第二行"
|
||
|
||
|
||
def test_clean_ocr_text_keeps_inline_backticks() -> None:
|
||
"""行内反引号属于正文内容,不能被误删。"""
|
||
# 数据:正文含行内反引号。
|
||
raw = "输入 `git status` 查看状态"
|
||
|
||
# 测试过程与验证结果
|
||
assert _clean_ocr_text(raw) == "输入 `git status` 查看状态"
|
||
|
||
|
||
def test_extract_gettext_takes_first_tag_to_avoid_loops() -> None:
|
||
"""多个 <gettext> 标签时只取第一个(防重复循环)。"""
|
||
# 数据:两个标签,模型循环输出了两次。
|
||
raw = "<gettext>正确文本</gettext><gettext>正确文本</gettext>"
|
||
|
||
# 测试过程与验证结果
|
||
assert _extract_gettext(raw) == "正确文本"
|
||
|
||
|
||
def test_extract_gettext_falls_back_to_raw_without_tags() -> None:
|
||
"""模型未按格式输出标签时回退原文。"""
|
||
# 数据:无标签的纯文本。
|
||
# 测试过程与验证结果
|
||
assert _extract_gettext("没有标签的文本") == "没有标签的文本"
|
||
|
||
|
||
def test_truncate_at_stop_cuts_at_earliest_sequence() -> None:
|
||
"""在最早命中的终止序列处截断(换行优先于"答")。"""
|
||
# 数据:文本中包含换行与"答"。
|
||
raw = "识别结果\n答:多余内容"
|
||
|
||
# 测试过程与验证结果
|
||
assert _truncate_at_stop(raw) == "识别结果"
|
||
|
||
|
||
def test_truncate_at_stop_returns_raw_when_no_match() -> None:
|
||
"""未命中终止序列时原样返回。"""
|
||
# 数据:不含任何终止序列。
|
||
# 测试过程与验证结果
|
||
assert _truncate_at_stop("普通文本") == "普通文本"
|
||
|
||
|
||
def test_stop_sequence_includes_newline_first() -> None:
|
||
"""终止序列以换行为首(输出首个换行即停),并包含"答"类重复模式。"""
|
||
# 数据:模块常量。
|
||
# 测试过程与验证结果
|
||
assert STOP_SEQUENCE[0] == "\n"
|
||
assert "答" in STOP_SEQUENCE
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 流式读取
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_consume_stream_joins_chunks() -> None:
|
||
"""多块流式内容按序拼接返回。"""
|
||
# 数据:三个内容块 + done 行。
|
||
response = _FakeResponse(_stream_lines("前", "中", "后"))
|
||
|
||
# 测试过程
|
||
raw = _consume_stream(response, deadline=1e18)
|
||
|
||
# 验证结果
|
||
assert raw == "前中后"
|
||
|
||
|
||
def test_consume_stream_stops_at_stop_sequence() -> None:
|
||
"""命中终止序列后不再读取后续块(防止无限循环输出)。"""
|
||
# 数据:第一块已含换行,第二块是循环垃圾。
|
||
response = _FakeResponse(_stream_lines("识别结果\n", "循环垃圾"))
|
||
|
||
# 测试过程
|
||
raw = _consume_stream(response, deadline=1e18)
|
||
|
||
# 验证结果:只取到第一个块(含换行)。
|
||
assert raw == "识别结果\n"
|
||
assert "循环垃圾" not in raw
|
||
|
||
|
||
def test_consume_stream_breaks_on_done_without_message() -> None:
|
||
"""done 行不带 message 时视为正常结束,不报错。"""
|
||
# 数据:仅 done 行。
|
||
response = _FakeResponse([json.dumps({"done": True}).encode("utf-8")])
|
||
|
||
# 测试过程与验证结果
|
||
assert _consume_stream(response, deadline=1e18) == ""
|
||
|
||
|
||
def test_consume_stream_raises_on_missing_message() -> None:
|
||
"""既无 message 又无 done 的块属于格式错误,明确报错。"""
|
||
# 数据:一个非法块。
|
||
response = _FakeResponse([json.dumps({"unexpected": 1}).encode("utf-8")])
|
||
|
||
# 测试过程与验证结果
|
||
with pytest.raises(KeyError):
|
||
_consume_stream(response, deadline=1e18)
|
||
|
||
|
||
def test_consume_stream_raises_on_deadline() -> None:
|
||
"""超过整体截止时间立即终止(每次调用 5 秒上限)。"""
|
||
# 数据:截止时间已过。
|
||
response = _FakeResponse(_stream_lines("内容", done=False))
|
||
|
||
# 测试过程与验证结果
|
||
with pytest.raises(TimeoutError):
|
||
_consume_stream(response, deadline=0.0)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# invoke:请求体结构与产物
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def test_invoke_sends_system_prompt_and_image_only_user_message(monkeypatch, tmp_path: Path) -> None:
|
||
"""请求体结构与 glm-ocr 期望一致:指令在 system,user 只带图片。"""
|
||
# 数据:捕获真实构造的 HTTP 请求体。
|
||
captured: dict = {}
|
||
|
||
def fake_urlopen(http_request, timeout=None):
|
||
captured["url"] = http_request.full_url
|
||
captured["body"] = json.loads(http_request.data.decode("utf-8"))
|
||
captured["timeout"] = timeout
|
||
return _FakeResponse(_stream_lines("识别文本"))
|
||
|
||
monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
|
||
|
||
# 测试过程
|
||
response = invoke(_request(tmp_path, TEXT_IMAGE, ollama_host="http://example:1"))
|
||
|
||
# 验证结果:URL、流式、系统提示词、user 只带 images、stop 与采样参数。
|
||
assert response.status == "completed", response.error
|
||
assert captured["url"] == "http://example:1/api/chat"
|
||
body = captured["body"]
|
||
assert body["stream"] is True
|
||
assert body["messages"][0]["role"] == "system"
|
||
assert body["messages"][0]["content"]
|
||
assert body["messages"][1]["content"] == ""
|
||
assert len(body["messages"][1]["images"]) == 1
|
||
assert body["stop"] == STOP_SEQUENCE
|
||
assert body["options"]["num_predict"] == 256
|
||
assert "keep_alive" in body
|
||
|
||
|
||
def test_invoke_writes_ocr_artifact(monkeypatch, tmp_path: Path) -> None:
|
||
"""识别结果写入 ocr.txt 并返回 text / text_uri 两个输出。
|
||
|
||
真实调用约定:模型输出首个换行即停(stop=\"\\n\"),因此这里用单行输出;
|
||
围栏清洗由 _clean_ocr_text 的独立用例覆盖。
|
||
"""
|
||
# 数据:假响应返回单行识别结果 + 后续循环垃圾。
|
||
monkeypatch.setattr(
|
||
urllib.request, "urlopen",
|
||
lambda *a, **k: _FakeResponse(_stream_lines("你好\n", "循环垃圾")),
|
||
)
|
||
|
||
# 测试过程
|
||
response = invoke(_request(tmp_path, TEXT_IMAGE))
|
||
|
||
# 验证结果:产物只含首个换行之前的内容。
|
||
assert response.status == "completed"
|
||
assert response.outputs["text"] == "你好"
|
||
artifact = Path(response.outputs["text_uri"])
|
||
assert artifact.read_text(encoding="utf-8").strip() == "你好"
|
||
|
||
|
||
def test_invoke_options_are_overridable(monkeypatch, tmp_path: Path) -> None:
|
||
"""temperature / repeat_penalty / num_predict 可被参数覆盖。"""
|
||
# 数据:捕获请求体。
|
||
captured: dict = {}
|
||
|
||
def fake_urlopen(http_request, timeout=None):
|
||
captured["body"] = json.loads(http_request.data.decode("utf-8"))
|
||
return _FakeResponse(_stream_lines("x"))
|
||
|
||
monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
|
||
|
||
# 测试过程
|
||
invoke(_request(
|
||
tmp_path, TEXT_IMAGE,
|
||
temperature=0.9, repeat_penalty=1.5, num_predict=64, timeout_seconds=2,
|
||
))
|
||
|
||
# 验证结果
|
||
options = captured["body"]["options"]
|
||
assert options["temperature"] == 0.9
|
||
assert options["repeat_penalty"] == 1.5
|
||
assert options["num_predict"] == 64
|
||
|
||
|
||
def test_invoke_fails_without_image_uri(tmp_path: Path) -> None:
|
||
"""缺少 image_uri 时返回 failed。"""
|
||
# 数据:空输入。
|
||
# 测试过程
|
||
response = invoke(_request(tmp_path, None))
|
||
|
||
# 验证结果
|
||
assert response.status == "failed"
|
||
assert "image_uri" in (response.error or "")
|
||
|
||
|
||
def test_invoke_fails_when_image_missing(tmp_path: Path) -> None:
|
||
"""图片文件不存在时提前失败,不发网络请求。"""
|
||
# 数据:不存在的图片路径。
|
||
# 测试过程
|
||
response = invoke(_request(tmp_path, tmp_path / "nope.png"))
|
||
|
||
# 验证结果
|
||
assert response.status == "failed"
|
||
assert "not found" in (response.error or "")
|
||
|
||
|
||
def test_invoke_fails_on_network_error(monkeypatch, tmp_path: Path) -> None:
|
||
"""连接失败时返回 failed 并带上错误信息(不抛异常到调度器外)。"""
|
||
# 数据:urlopen 抛 URLError。
|
||
def fake_urlopen(*a, **k):
|
||
raise urllib.error.URLError("connection refused")
|
||
|
||
monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
|
||
|
||
# 测试过程
|
||
response = invoke(_request(tmp_path, TEXT_IMAGE))
|
||
|
||
# 验证结果
|
||
assert response.status == "failed"
|
||
assert "connection refused" in (response.error or "")
|
||
|
||
|
||
def test_invoke_fails_on_timeout(monkeypatch, tmp_path: Path) -> None:
|
||
"""超时(socket.timeout)时返回 failed,而不是长时间挂起。"""
|
||
# 数据:urlopen 抛 socket.timeout。
|
||
def fake_urlopen(*a, **k):
|
||
raise socket.timeout("timed out")
|
||
|
||
monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
|
||
|
||
# 测试过程
|
||
response = invoke(_request(tmp_path, TEXT_IMAGE))
|
||
|
||
# 验证结果
|
||
assert response.status == "failed"
|
||
|
||
|
||
def test_invoke_fails_on_malformed_stream(monkeypatch, tmp_path: Path) -> None:
|
||
"""流式块格式错误时返回 failed。"""
|
||
# 数据:非法块。
|
||
monkeypatch.setattr(
|
||
urllib.request, "urlopen",
|
||
lambda *a, **k: _FakeResponse([json.dumps({"x": 1}).encode("utf-8")]),
|
||
)
|
||
|
||
# 测试过程
|
||
response = invoke(_request(tmp_path, TEXT_IMAGE))
|
||
|
||
# 验证结果
|
||
assert response.status == "failed"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 真实模型集成
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _ollama_reachable(host: str = "http://192.168.123.70:11434", model: str = "glm-ocr:latest") -> bool:
|
||
"""探测真实 Ollama 服务与目标模型是否可用(缺失即跳过集成测试)。"""
|
||
try:
|
||
request = urllib.request.Request(
|
||
f"{host}/api/show",
|
||
data=json.dumps({"model": model}).encode("utf-8"),
|
||
headers={"Content-Type": "application/json"},
|
||
method="POST",
|
||
)
|
||
with urllib.request.urlopen(request, timeout=5) as resp:
|
||
return resp.status == 200
|
||
except (urllib.error.URLError, OSError):
|
||
return False
|
||
|
||
|
||
@pytest.mark.integration
|
||
def test_vlm_ocr_real_model_on_real_subtitle_image(tmp_path: Path) -> None:
|
||
"""真实 glm-ocr + 真实字幕截图:应识别出期望字幕文本,且无围栏垃圾。"""
|
||
# 数据:模块 data/ 下的真实字幕截图。
|
||
if not REAL_SUBTITLE_IMAGE.is_file():
|
||
pytest.skip(f"缺少测试资产 {REAL_SUBTITLE_IMAGE},跳过")
|
||
if not _ollama_reachable():
|
||
pytest.skip("Ollama 服务或 glm-ocr 模型不可用,跳过真实模型集成测试")
|
||
|
||
# 测试过程
|
||
response = invoke(_request(
|
||
tmp_path, REAL_SUBTITLE_IMAGE,
|
||
model="glm-ocr:latest", ollama_host="http://192.168.123.70:11434",
|
||
))
|
||
|
||
# 验证结果:包含期望文本(允许尾部重复循环被截断),且无围栏。
|
||
assert response.status == "completed", response.error
|
||
text = response.outputs["text"]
|
||
assert EXPECTED_REAL_TEXT in text, f"未识别出期望字幕:{text[:200]}"
|
||
assert "```" not in text
|
||
|
||
|
||
@pytest.mark.integration
|
||
def test_vlm_ocr_real_model_on_no_text_image(tmp_path: Path) -> None:
|
||
"""真实无文字帧:模型不应输出画面描述文字(OCR 只提文字不做描述)。"""
|
||
# 数据:无文字测试图。
|
||
if not NOTEXT_IMAGE.is_file():
|
||
pytest.skip(f"缺少测试资产 {NOTEXT_IMAGE},跳过")
|
||
if not _ollama_reachable():
|
||
pytest.skip("Ollama 服务或 glm-ocr 模型不可用,跳过真实模型集成测试")
|
||
|
||
# 测试过程
|
||
response = invoke(_request(
|
||
tmp_path, NOTEXT_IMAGE,
|
||
model="glm-ocr:latest", ollama_host="http://192.168.123.70:11434",
|
||
))
|
||
|
||
# 验证结果:要么成功且文本很短(无文字),要么失败——都不应出现长描述。
|
||
if response.status == "completed":
|
||
assert len(response.outputs["text"].strip()) <= 40
|