"""nodes/subtitle_ocr.py 的模块级测试（数据 → 测试过程 → 验证结果）。 被测模块：`nodes/subtitle_ocr.py`（逐帧 OCR → 合并 → 汇总 SRT，含并发与 断点续跑），可独立调用。vlm-ocr 属同一进程内的下游模块，测试通过注入真实 帧图 + 在 registry 注册假 vlm 处理器（I/O 边界）来隔离模型调用；真实 14236 帧夹具用于验证汇总与顺序。 """ from __future__ import annotations import json from pathlib import Path import pytest from nodes.adaptive_pool import AdaptiveThreadPool from nodes.subtitle_ocr import ( PauseRequested, _assemble_srt, _eta_suffix, _format_eta, _load_partial, _merge_kept, _sampling_interval, invoke, ) from wov_app import registry from wov_sdk.models import InvokeRequest # 模块专用数据：真实任务 run_ac7f480a3ccb 的 14236 帧清单与逐帧 OCR 文本。 DATA_DIR = Path(__file__).resolve().parent / "data" FULL_MANIFEST = DATA_DIR / "frames_manifest_full.json" FULL_TEXTS = DATA_DIR / "ocr_frames_full.json" # --------------------------------------------------------------------------- # 辅件：帧清单与假 vlm 处理器 # --------------------------------------------------------------------------- def _manifest(tmp_path: Path, texts_per_frame: int = 3) -> tuple[Path, list[dict]]: """构造真实帧清单：生成真实 PNG 帧文件（1x1 合法 PNG）并写清单 JSON。 返回 (清单路径, 清单数据)。时间轴按 2 秒间隔，与真实抽帧一致。 """ frames_dir = tmp_path / "frames" frames_dir.mkdir(parents=True, exist_ok=True) # 最小合法 PNG（1x1 透明像素），供下游假处理器读取真实文件。 png = bytes.fromhex( "89504e470d0a1a0a0000000d49484452000000010000000108060000001f15c489" "0000000a49444154789c6300010000050001od".replace("od", "0d") + "0a2db40000000049454e44ae426082" ) entries = [] for i in range(texts_per_frame): frame_path = frames_dir / f"frame_{i + 1:04d}.png" frame_path.write_bytes(png) entries.append({"time": round(i * 2.0, 3), "image_uri": str(frame_path)}) manifest_path = tmp_path / "frames.json" manifest_path.write_text(json.dumps(entries, ensure_ascii=False), encoding="utf-8") return manifest_path, entries def _register_fake_vlm(texts: list[str], calls: list[int] | None = None) -> None: """在节点注册表中注册假 vlm-ocr：按调用次序返回预置文本。 这是允许的 I/O 边界替身（模型推理不进入单元测试）；使用仓库真实的 vlm 清单（manifests/vlm.json）注册，保证注册表校验路径也被执行， 返回结构为真实 InvokeResponse（含 text 字段）。 """ from pathlib import Path as _Path from wov_sdk.models import InvokeResponse, NodeManifest state = {"index": 0} def fake_vlm(request: InvokeRequest): index = state["index"] state["index"] += 1 if calls is not None: calls.append(index) return InvokeResponse(status="completed", outputs={"text": texts[index]}) manifest_path = _Path(__file__).resolve().parents[3] / "manifests" / "vlm.json" registry.register(NodeManifest.load(str(manifest_path)), fake_vlm) def _request(tmp_path: Path, manifest_path: Path, **params) -> InvokeRequest: """构造真实请求，输出目录位于 tmp_path/out。""" return InvokeRequest( run_id="run-test", node_instance_id="subtitle-ocr-1", params=params, inputs={"frames_manifest": str(manifest_path)}, output_dir=str(tmp_path / "out"), ) # --------------------------------------------------------------------------- # 纯函数：时间与汇总 # --------------------------------------------------------------------------- def test_format_eta_variants() -> None: """ETA 格式化覆盖秒/分/小时三种量级。""" # 数据：45 秒、34 分 13 秒、2 小时 5 分。 # 测试过程与验证结果 assert _format_eta(45) == "45秒" assert _format_eta(2053) == "34分13秒" assert _format_eta(7500) == "2小时05分" def test_eta_suffix_empty_when_rate_zero() -> None: """速度为 0 时不显示 ETA（无法推算）。""" # 数据：速度为 0 与正常速度。 # 测试过程与验证结果 assert _eta_suffix(10, 100, 0.0) == "" assert "预计剩余" in _eta_suffix(50, 100, 1.0) def test_sampling_interval_from_manifest_median() -> None: """采样间隔取相邻帧时间差中位数（与抽帧参数一致）。""" # 数据：0.5 秒间隔的帧清单。 manifest = [{"time": round(i * 0.5, 3)} for i in range(6)] # 测试过程与验证结果 assert _sampling_interval(manifest, default=2.0) == 0.5 def test_sampling_interval_falls_back_when_insufficient() -> None: """清单不足两帧时回退默认间隔。""" # 数据：单帧与空清单。 # 测试过程与验证结果 assert _sampling_interval([{"time": 0.0}], default=1.5) == 1.5 assert _sampling_interval([], default=1.5) == 1.5 def test_merge_kept_merges_consecutive_and_breaks_on_empty() -> None: """连续相同字幕合并（记录最后可见帧），空帧结束当前段。""" # 数据：A A 空 A —— 第一段两条 A，空帧后是新的一段。 manifest = [{"time": t} for t in (0.0, 2.0, 4.0, 6.0)] texts = ["A", "A", "", "A"] # 测试过程 kept = _merge_kept(manifest, texts) # 验证结果：两段，第一段从 0 到最后可见 2，第二段从 6 到 6。 assert kept == [(0.0, 2.0, "A"), (6.0, 6.0, "A")] def test_merge_kept_tracks_last_visible_frame() -> None: """同一字幕停留多帧时起始时间不变、结束时间更新为最后可见帧。""" # 数据：三条相同字幕。 manifest = [{"time": t} for t in (0.0, 2.0, 4.0)] # 测试过程 kept = _merge_kept(manifest, ["字幕", "字幕", "字幕"]) # 验证结果 assert kept == [(0.0, 4.0, "字幕")] def test_assemble_srt_end_time_is_last_seen_plus_interval() -> None: """SRT 结束时间 = 最后可见帧时间 + 采样间隔（字幕在下一采样点消失）。""" # 数据：一段字幕，采样间隔 0.5。 kept = [(1.0, 2.0, "台词")] # 测试过程 lines = _assemble_srt(kept, 0.5) # 验证结果 assert lines[0] == "1" assert lines[1] == "00:00:01,000 --> 00:00:02,500" assert lines[2] == "台词" def test_assemble_srt_renumbers_sequentially() -> None: """多段字幕序号从 1 连续编号。""" # 数据：两段字幕。 kept = [(0.0, 1.0, "甲"), (5.0, 6.0, "乙")] # 测试过程 lines = _assemble_srt(kept, 1.0) # 验证结果：序号 1、2。 text = "\n".join(lines) assert "\n1\n" in f"\n{text}" or text.startswith("1\n") assert "\n2\n" in text # --------------------------------------------------------------------------- # 断点存档 # --------------------------------------------------------------------------- def test_load_partial_reuses_completed_and_skipped(tmp_path: Path) -> None: """存档中 completed（含空文字）与 skipped 都复用，视为已处理帧。""" # 数据：一份含成功空帧、成功文本、跳过、以及旧版空串的存档。 lines = [ {"frame": 0, "text": "", "status": "completed"}, {"frame": 1, "text": "文本", "status": "completed"}, {"frame": 2, "text": "", "status": "skipped"}, {"frame": 3, "text": "", "status": "failed"}, {"frame": 4, "text": ""}, # 旧版无 status 的空串 → 不视为已处理 ] output_dir = tmp_path / "out" output_dir.mkdir() (output_dir / "ocr_partial.jsonl").write_text( "\n".join(json.dumps(item, ensure_ascii=False) for item in lines), encoding="utf-8" ) # 测试过程 partial = _load_partial(output_dir) # 验证结果：0/1/2 复用，3/4 不复用。 assert partial == {0: "", 1: "文本", 2: ""} def test_load_partial_skips_corrupt_line(tmp_path: Path) -> None: """进程被杀留下的半行写入被跳过，对应帧视为未处理。""" # 数据：合法行 + 截断行。 output_dir = tmp_path / "out" output_dir.mkdir() (output_dir / "ocr_partial.jsonl").write_text( json.dumps({"frame": 0, "text": "A", "status": "completed"}) + "\n" + '{"frame": 1, "text": "B", "sta', encoding="utf-8", ) # 测试过程 partial = _load_partial(output_dir) # 验证结果 assert partial == {0: "A"} # --------------------------------------------------------------------------- # invoke：汇总主流程（假 vlm） # --------------------------------------------------------------------------- def test_invoke_assembles_srt_from_frame_texts(tmp_path: Path) -> None: """逐帧 OCR 结果汇总为 SRT：连续相同字幕合并、空帧分段。""" # 数据：A A 空 的帧序列。 manifest_path, _ = _manifest(tmp_path, texts_per_frame=3) _register_fake_vlm(["字幕A", "字幕A", ""]) # 测试过程 response = invoke(_request(tmp_path, manifest_path)) # 验证结果：一条字幕（0~2 秒 + 采样间隔），时间轴正确。 assert response.status == "completed", response.error content = Path(response.outputs["srt_uri"]).read_text(encoding="utf-8") assert "字幕A" in content assert response.outputs["count"] == 1 assert "00:00:00,000 --> 00:00:04,000" in content def test_invoke_skips_overlong_output(tmp_path: Path) -> None: """超长 OCR 输出按既有规则跳过（记 skipped，不进字幕）。""" # 数据：一帧超长输出 + 一帧正常。 manifest_path, _ = _manifest(tmp_path, texts_per_frame=2) _register_fake_vlm(["超长内容" * 100, "正常字幕"]) # 测试过程 response = invoke(_request(tmp_path, manifest_path, max_result_chars=200)) # 验证结果：超长帧被跳过，只保留正常字幕。 assert response.status == "completed", response.error content = Path(response.outputs["srt_uri"]).read_text(encoding="utf-8") assert "正常字幕" in content assert "超长内容" not in content def test_invoke_resumes_from_checkpoint(tmp_path: Path) -> None: """断点续跑：已存档帧不再调用 vlm，只处理剩余帧。""" # 数据：3 帧，其中第 0 帧已有存档。 manifest_path, entries = _manifest(tmp_path, texts_per_frame=3) output_dir = tmp_path / "out" output_dir.mkdir() (output_dir / "ocr_partial.jsonl").write_text( json.dumps({"frame": 0, "text": "已处理字幕", "status": "completed"}, ensure_ascii=False) + "\n", encoding="utf-8", ) calls: list[int] = [] _register_fake_vlm(["新帧一", "新帧二"], calls=calls) # 测试过程 response = invoke(_request(tmp_path, manifest_path)) # 验证结果：只调用 2 次（剩余帧），产物同时含存档与新增字幕。 assert response.status == "completed", response.error assert len(calls) == 2 content = Path(response.outputs["srt_uri"]).read_text(encoding="utf-8") assert "已处理字幕" in content assert "新帧一" in content def test_invoke_all_frames_cached_makes_no_vlm_call(tmp_path: Path) -> None: """全部帧已存档时不再调用 vlm（幂等重跑）。""" # 数据：2 帧全部已存档。 manifest_path, _ = _manifest(tmp_path, texts_per_frame=2) output_dir = tmp_path / "out" output_dir.mkdir() (output_dir / "ocr_partial.jsonl").write_text( "\n".join([ json.dumps({"frame": 0, "text": "甲", "status": "completed"}, ensure_ascii=False), json.dumps({"frame": 1, "text": "乙", "status": "completed"}, ensure_ascii=False), ]) + "\n", encoding="utf-8", ) calls: list[int] = [] _register_fake_vlm([], calls=calls) # 测试过程 response = invoke(_request(tmp_path, manifest_path)) # 验证结果：零调用且产物完整。 assert response.status == "completed", response.error assert calls == [] content = Path(response.outputs["srt_uri"]).read_text(encoding="utf-8") assert "甲" in content and "乙" in content def test_invoke_writes_checkpoint_per_frame(tmp_path: Path) -> None: """每帧完成后立即写断点存档（进程被杀后可续跑）。""" # 数据：2 帧。 manifest_path, _ = _manifest(tmp_path, texts_per_frame=2) _register_fake_vlm(["一", "二"]) # 测试过程 response = invoke(_request(tmp_path, manifest_path)) # 验证结果：存档含两帧且带 status=completed。 assert response.status == "completed" archived = (tmp_path / "out" / "ocr_partial.jsonl").read_text(encoding="utf-8").splitlines() records = [json.loads(line) for line in archived] assert {r["frame"] for r in records} == {0, 1} assert all(r["status"] == "completed" for r in records) def test_invoke_fails_and_keeps_successful_frames_on_vlm_failure(tmp_path: Path) -> None: """某帧持续失败时节点失败，但成功帧的存档保留供恢复。""" # 数据：3 帧，vlm 对第 1 帧始终失败。 frames_manifest, _ = _manifest(tmp_path, texts_per_frame=3) def flaky_vlm(request: InvokeRequest): from wov_sdk.models import InvokeResponse image = Path(request.inputs["image_uri"]).name if image == "frame_0002.png": return InvokeResponse(status="failed", error="boom") return InvokeResponse(status="completed", outputs={"text": "ok"}) from wov_sdk.models import NodeManifest # 注册假 vlm-ocr（使用真实清单文件，变量名与帧清单区分开）。 vlm_manifest_path = Path(__file__).resolve().parents[3] / "manifests" / "vlm.json" registry.register(NodeManifest.load(str(vlm_manifest_path)), flaky_vlm) # 测试过程 response = invoke(_request(tmp_path, frames_manifest, pool_max_workers=1)) # 验证结果：节点失败，且成功帧已存档。 assert response.status == "failed" archived = (tmp_path / "out" / "ocr_partial.jsonl").read_text(encoding="utf-8").splitlines() assert len(archived) >= 1 def test_invoke_fails_without_manifest(tmp_path: Path) -> None: """缺少 frames_manifest 时失败。""" # 数据：空输入。 request = InvokeRequest( run_id="r", node_instance_id="n", params={}, inputs={}, output_dir=str(tmp_path) ) # 测试过程 response = invoke(request) # 验证结果 assert response.status == "failed" assert "frames_manifest" in (response.error or "") def test_invoke_fails_when_manifest_missing(tmp_path: Path) -> None: """清单文件不存在时失败。""" # 数据：不存在的路径。 # 测试过程 response = invoke(_request(tmp_path, tmp_path / "nope.json")) # 验证结果 assert response.status == "failed" assert "not found" in (response.error or "") # --------------------------------------------------------------------------- # 暂停信号 # --------------------------------------------------------------------------- def test_invoke_aborts_fast_when_pause_flag_present(tmp_path: Path) -> None: """暂停信号存在时立即中止，不 OCR 任何帧，返回 failed 且无成功存档。""" # 数据：帧清单 + run 根目录下的 paused.flag（output_dir 的上两级）。 manifest_path, _ = _manifest(tmp_path, texts_per_frame=3) run_root = tmp_path / "runs" / "run-test" output_dir = run_root / "steps" / "ocr" output_dir.mkdir(parents=True) (run_root / "paused.flag").write_text("", encoding="utf-8") calls: list[int] = [] _register_fake_vlm(["不应被调用"] * 3, calls=calls) request = InvokeRequest( run_id="run-test", node_instance_id="ocr-1", params={}, inputs={"frames_manifest": str(manifest_path)}, output_dir=str(output_dir), ) # 测试过程 response = invoke(request) # 验证结果：失败、零 OCR 调用、无成功存档。 assert response.status == "failed" assert "暂停" in (response.error or "") assert calls == [] assert not (output_dir / "ocr_partial.jsonl").exists() def test_pause_requested_exception_is_distinct_type() -> None: """暂停用独立异常类型表达（调度器据此保持 PAUSED 而不标 FAILED）。""" # 数据：异常类。 # 测试过程与验证结果 assert issubclass(PauseRequested, Exception) # --------------------------------------------------------------------------- # 真实全量帧数据回归 # --------------------------------------------------------------------------- @pytest.mark.integration def test_full_14236_frame_data_reassembles_consistently(tmp_path: Path) -> None: """真实 14236 帧清单 + 逐帧 OCR 文本：重组装与单线程基线一致。 数据来源：真实任务 run_ac7f480a3ccb（见 docs/testing.md）。 """ # 数据：真实清单与真实逐帧文本。 if not FULL_MANIFEST.is_file() or not FULL_TEXTS.is_file(): pytest.skip("缺少真实 14236 帧夹具，跳过") # 测试过程：用真实数据直接调用汇总逻辑（不调模型）。 manifest = json.loads(FULL_MANIFEST.read_text(encoding="utf-8")) texts = json.loads(FULL_TEXTS.read_text(encoding="utf-8")) if isinstance(texts, dict): texts = [texts.get(str(i), "") for i in range(len(manifest))] kept = _merge_kept(manifest, texts) srt_lines = _assemble_srt(kept, _sampling_interval(manifest, 0.5)) # 验证结果：条数稳定、时间轴单调不减、SRT 结构合法。 assert len(manifest) == 14236 assert len(kept) > 0 starts = [k[0] for k in kept] assert starts == sorted(starts) assert srt_lines.count("") == len(kept) # 每条一个空行分隔