Files
vrsub/tests/nodes/test_vlm/test_ocr.py
cat-shark 8a715a8064 test: 按模块重写测试代码,删除旧平铺结构
按"测试规则"重写 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。
2026-09-13 15:40:56 +08:00

411 lines
14 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""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 期望一致:指令在 systemuser 只带图片。"""
# 数据:捕获真实构造的 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