"""抽帧与字幕 OCR 节点单元测试。 frame-extract 用 testdata 真实视频抽帧+裁切+720p 压缩;subtitle-ocr 的 OCR 网络调用(vlm-ocr)按 I/O 边界 mock,但喂给它的帧图片是真实的 testdata 资产。 应用层不再加工模型输出文本,只做长度上限校验(超长报错跳过)。 """ import json from pathlib import Path from wov_sdk.models import InvokeRequest, InvokeResponse from nodes.frame_extract import invoke as frame_invoke from nodes.subtitle_ocr import _assemble_srt from nodes.subtitle_ocr import invoke as ocr_invoke # 单体根目录:tests/ 的上一级。 WORKSPACE = Path(__file__).resolve().parent.parent TESTDATA = WORKSPACE / "testdata" # 10s 测试视频:SUB 001 在 1-4s、SUB 002 在 6-9s。 VIDEO = TESTDATA / "subtitle_10s.mp4" TEXT_IMG = TESTDATA / "ocr_text.png" def _png_size(path: Path) -> tuple[int, int]: """从 PNG 头读取宽高(真实图片尺寸断言)。""" data = path.read_bytes() assert data[:8] == b"\x89PNG\r\n\x1a\n", "not a real png" width = int.from_bytes(data[16:20], "big") height = int.from_bytes(data[20:24], "big") return width, height def _frame_request(tmp_path, video=VIDEO, **params) -> InvokeRequest: """构造 frame-extract 调用请求。""" return InvokeRequest( run_id="run_fx", node_instance_id="", inputs={"video_uri": str(video)}, params=params, output_dir=str(tmp_path / "out"), ) # --------------------------------------------------------------------------- # frame-extract:抽帧 + 裁切 + 720p 压缩 # --------------------------------------------------------------------------- def test_frame_extract_crop_and_manifest(tmp_path) -> None: """真实视频抽帧:裁切下半 50% 后帧尺寸为 1280x360,清单时间轴正确。""" response = frame_invoke( _frame_request(tmp_path, interval_seconds=1, crop=[0, 0.5, 1, 0.5]) ) assert response.status == "completed", response.error manifest = json.loads(Path(response.outputs["frames_manifest"]).read_text(encoding="utf-8")) assert len(manifest) >= 9 assert [round(item["time"], 3) for item in manifest] == [ round(i * 1.0, 3) for i in range(len(manifest)) ] first = Path(manifest[0]["image_uri"]) assert first.is_file() # 1280x360 已在 720p 内,压缩不改变尺寸。 assert _png_size(first) == (1280, 360) def test_frame_extract_default_params(tmp_path) -> None: """未指定参数时使用默认值:抽帧间隔 0.5 秒 + 默认底部裁切区域。""" response = frame_invoke(_frame_request(tmp_path)) assert response.status == "completed", response.error manifest = json.loads(Path(response.outputs["frames_manifest"]).read_text(encoding="utf-8")) assert manifest # 默认间隔 0.5s:25fps 下 step=round(12.5)=12(银行家舍入), # 帧时间按 step/fps=12/25=0.48s 步进(帧号精确,采样周期由帧量化决定)。 assert [round(item["time"], 3) for item in manifest] == [ round(i * 12 / 25, 3) for i in range(len(manifest)) ] def test_frame_extract_missing_video(tmp_path) -> None: """缺少 video_uri 时返回失败。""" response = frame_invoke( InvokeRequest( run_id="r", node_instance_id="", inputs={}, output_dir=str(tmp_path) ) ) assert response.status == "failed" def test_frame_extract_bad_crop(tmp_path) -> None: """crop 比例越界(超出画面)时返回失败。""" assert frame_invoke(_frame_request(tmp_path, crop=[0, 0.5, 1, 1.5])).status == "failed" assert frame_invoke(_frame_request(tmp_path, crop=[-0.1, 0, 1, 0.5])).status == "failed" assert frame_invoke(_frame_request(tmp_path, crop="abc")).status == "failed" assert frame_invoke(_frame_request(tmp_path, crop=[0, 0.5, 1])).status == "failed" # 各值域合法但 x+w 越出画面。 assert frame_invoke(_frame_request(tmp_path, crop=[0.6, 0, 0.5, 0.3])).status == "failed" def test_frame_extract_video_missing_file(tmp_path) -> None: """video_uri 指向不存在的文件时返回失败。""" response = frame_invoke(_frame_request(tmp_path, video=tmp_path / "none.mp4")) assert response.status == "failed" assert "not found" in response.error def test_frame_extract_bad_interval(tmp_path) -> None: """间隔 <= 0 时返回失败。""" response = frame_invoke(_frame_request(tmp_path, interval_seconds=0)) assert response.status == "failed" def test_frame_extract_ffmpeg_fails(monkeypatch, tmp_path) -> None: """ffmpeg 抽帧失败时透传错误。""" import subprocess as sp monkeypatch.setattr("nodes.frame_extract._video_size", lambda *a, **k: (1280, 720)) monkeypatch.setattr("nodes.frame_extract._video_fps", lambda *a, **k: 25.0) monkeypatch.setattr("nodes.frame_extract._video_duration", lambda *a, **k: 10.0) monkeypatch.setattr( "nodes.frame_extract.subprocess.run", lambda *a, **k: sp.CompletedProcess([], 1, stderr="boom"), ) response = frame_invoke(_frame_request(tmp_path)) assert response.status == "failed" assert "boom" in response.error def test_video_size_unreadable(monkeypatch) -> None: """ffmpeg -i 输出不含视频流信息时返回 None。""" import subprocess as sp from nodes.frame_extract import _video_size monkeypatch.setattr( "nodes.frame_extract.subprocess.run", lambda *a, **k: sp.CompletedProcess([], 0, stderr="no video stream"), ) assert _video_size(Path("/tmp/x.mp4"), "ffmpeg") is None def test_frame_extract_video_size_unknown(monkeypatch, tmp_path) -> None: """无法读取视频分辨率时返回失败。""" monkeypatch.setattr("nodes.frame_extract._video_size", lambda *a, **k: None) response = frame_invoke(_frame_request(tmp_path)) assert response.status == "failed" assert "video size" in response.error def test_video_duration_unreadable(monkeypatch) -> None: """ffmpeg -i 输出缺少 Duration 时返回 None。""" import subprocess as sp from nodes.frame_extract import _video_duration monkeypatch.setattr( "nodes.frame_extract.subprocess.run", lambda *a, **k: sp.CompletedProcess([], 0, stderr="no duration info"), ) assert _video_duration(Path("/tmp/x.mp4"), "ffmpeg") is None class _FakeReader: """模拟 stderr 读取对象。""" def __init__(self, content: str = "") -> None: self._content = content def read(self) -> str: return self._content class FakePopen: """模拟 ffmpeg 进程:stdout 可迭代 -progress 行,可配置返回码与 stderr。""" def __init__(self, lines=(), returncode: int = 0, stderr: str = "") -> None: self.stdout = list(lines) self.stderr = _FakeReader(stderr) self.returncode = returncode def wait(self) -> int: return self.returncode def _fake_popen(lines=(), returncode: int = 0, stderr: str = ""): """构造替换 subprocess.Popen 的工厂函数。""" return lambda *a, **k: FakePopen(lines, returncode, stderr) def test_parse_progress_line() -> None: """-progress 行解析:frame=N 返回数值,其他行与非法值返回 None。""" from nodes.frame_extract import _parse_progress_line assert _parse_progress_line("frame=25\n") == 25 assert _parse_progress_line("progress=continue") is None assert _parse_progress_line("fps=25.0") is None assert _parse_progress_line("frame=abc") is None def test_frame_extract_progress_logging(monkeypatch, tmp_path) -> None: """ffmpeg -progress 的 frame=N 被解析并打印抽帧进度与速度。""" import logging monkeypatch.setattr("nodes.frame_extract._video_size", lambda *a, **k: (1280, 720)) monkeypatch.setattr("nodes.frame_extract._video_fps", lambda *a, **k: 25.0) monkeypatch.setattr("nodes.frame_extract._video_duration", lambda *a, **k: 10.0) monkeypatch.setattr( "nodes.frame_extract.subprocess.Popen", _fake_popen(lines=["fps=25.0", "frame=20", "progress=continue", "frame=50", "progress=continue"]), ) captured: list[str] = [] class CaptureHandler(logging.Handler): def emit(self, record): captured.append(record.getMessage()) logger = logging.getLogger("vrsub.frame-extract") logger.addHandler(CaptureHandler()) try: response = frame_invoke(_frame_request(tmp_path)) finally: logger.removeHandler(logger.handlers[-1]) assert response.status == "completed", response.error assert any("抽帧进度" in message and "帧/s" in message for message in captured) def test_frame_extract_ffmpeg_fails(monkeypatch, tmp_path) -> None: """ffmpeg 抽帧失败时透传错误。""" monkeypatch.setattr("nodes.frame_extract._video_size", lambda *a, **k: (1280, 720)) monkeypatch.setattr("nodes.frame_extract._video_fps", lambda *a, **k: 25.0) monkeypatch.setattr("nodes.frame_extract._video_duration", lambda *a, **k: 10.0) monkeypatch.setattr( "nodes.frame_extract.subprocess.Popen", _fake_popen(lines=["frame=1", "progress=end"], returncode=1, stderr="boom"), ) response = frame_invoke(_frame_request(tmp_path)) assert response.status == "failed" assert "boom" in response.error def test_frame_extract_duration_unknown(monkeypatch, tmp_path) -> None: """无法读取视频时长时返回失败。""" monkeypatch.setattr("nodes.frame_extract._video_size", lambda *a, **k: (1280, 720)) monkeypatch.setattr("nodes.frame_extract._video_duration", lambda *a, **k: None) response = frame_invoke(_frame_request(tmp_path)) assert response.status == "failed" assert "duration" in response.error def test_frame_extract_fps_unknown(monkeypatch, tmp_path) -> None: """无法读取视频帧率时返回失败。""" monkeypatch.setattr("nodes.frame_extract._video_size", lambda *a, **k: (1280, 720)) monkeypatch.setattr("nodes.frame_extract._video_duration", lambda *a, **k: 10.0) monkeypatch.setattr("nodes.frame_extract._video_fps", lambda *a, **k: None) response = frame_invoke(_frame_request(tmp_path)) assert response.status == "failed" assert "fps" in response.error def test_frame_step_conversion() -> None: """帧间隔换算:step=round(间隔秒×fps),至少为 1。""" from nodes.frame_extract import _frame_step assert _frame_step(fps=25.0, interval=0.2) == 5 assert _frame_step(fps=25.0, interval=1.0) == 25 assert _frame_step(fps=29.97, interval=1.0) == 30 # fps 很低时 step 也不会小于 1(每帧都取)。 assert _frame_step(fps=1.0, interval=0.2) == 1 def test_video_fps_parse(monkeypatch) -> None: """帧率解析:支持小数(29.97)与有理数(30000/1001)。""" import subprocess as sp from nodes.frame_extract import _video_fps monkeypatch.setattr( "nodes.frame_extract.subprocess.run", lambda *a, **k: sp.CompletedProcess( [], 0, stderr="Stream #0:0: Video: h264, 1280x720, 30000/1001 fps, 30000/1001 tbr" ), ) assert _video_fps(Path("/tmp/x.mp4"), "ffmpeg") == 30000 / 1001 monkeypatch.setattr( "nodes.frame_extract.subprocess.run", lambda *a, **k: sp.CompletedProcess( [], 0, stderr="Stream #0:0: Video: h264, 1280x720, 25 fps, 25 tbr" ), ) assert _video_fps(Path("/tmp/x.mp4"), "ffmpeg") == 25.0 monkeypatch.setattr( "nodes.frame_extract.subprocess.run", lambda *a, **k: sp.CompletedProcess([], 0, stderr="no video stream"), ) assert _video_fps(Path("/tmp/x.mp4"), "ffmpeg") is None def test_frame_extract_one_second_exact_frames(tmp_path) -> None: """真实视频 1s 间隔按秒 seek 精确抽帧:10s 视频应得 10 帧,时间 0..9。""" response = frame_invoke( _frame_request(tmp_path, interval_seconds=1, crop=[0, 0.7, 1, 0.3]) ) assert response.status == "completed", response.error manifest = json.loads(Path(response.outputs["frames_manifest"]).read_text(encoding="utf-8")) assert [round(item["time"], 3) for item in manifest] == [ round(i * 1.0, 3) for i in range(len(manifest)) ] assert len(manifest) == 10 def test_frame_files_read_order_matches_frame_number(tmp_path) -> None: """真实任务留存数据:帧文件按帧号数值排序读取,而非字典序。 回归用例:ffmpeg 的 %04d 编号在超过 9999 帧后自动扩为 5 位 (frame_10000.png 等),此时 sorted() 默认字典序会把 5 位编号排在 4 位编号之前(如 frame_10009 < frame_1009),导致帧号回退、manifest 时间与图像错位。testdata/frames_boundary/ 是 2026-08 真实任务 run_339ec7ee437f(14236 帧 / 2 小时视频)中跨越该边界的真实帧文件。 """ from nodes.frame_extract import _sorted_frame_files boundary = TESTDATA / "frames_boundary" files = _sorted_frame_files(boundary) # 从文件名解析帧号:读取顺序必须等于帧号数值递增序(无回退)。 nums = [int(p.stem.split("_", 1)[1]) for p in files] assert nums == sorted(nums) # 边界关键对:5 位编号必须排在 4 位编号之后,禁止字典序错位。 assert nums.index(10000) > nums.index(9999) assert nums.index(10009) > nums.index(1009) # --------------------------------------------------------------------------- # subtitle-ocr:OCR 循环 + 长度上限 + 合并 + SRT 组装 # --------------------------------------------------------------------------- def _frames_manifest(tmp_path, frame_specs) -> Path: """构造真实 frames.json;frame_specs=[(time, image_path), ...]。""" items = [{"time": time, "image_uri": str(image)} for time, image in frame_specs] path = tmp_path / "frames.json" path.write_text(json.dumps(items), encoding="utf-8") return path def test_assemble_srt_real_timeline() -> None: """SRT 组装:起始=帧时间,结束=最后可见帧时间+采样间隔。""" lines = _assemble_srt([(0.0, 6.0, "A"), (6.0, 8.0, "B")], interval_seconds=2.0) text = "\n".join(lines) assert text.startswith("1\n") # A 最后可见帧 6.0 + 间隔 2.0 = 8.0(而非下一条字幕的出现时间)。 assert "00:00:00,000 --> 00:00:08,000" in text assert "00:00:06,000 --> 00:00:10,000" in text def test_sampling_interval_from_manifest() -> None: """采样间隔从帧清单时间轴推导:均匀间隔取相邻差,退化清单回退默认值。""" from nodes.subtitle_ocr import _sampling_interval manifest = [{"time": i * 0.2, "image_uri": f"f{i}.png"} for i in range(10)] assert _sampling_interval(manifest, 2.0) == 0.2 # 单帧(无法算差)与异常时间序:回退默认值。 assert _sampling_interval([{"time": 0.0, "image_uri": "f0.png"}], 2.0) == 2.0 assert _sampling_interval( [{"time": 0.0}, {"time": 0.0}, {"time": 0.2}], 2.0 ) == 0.2 def test_ocr_merges_consecutive_same_text(monkeypatch, tmp_path) -> None: """连续帧相同字幕合并为一条;消失时间=最后可见帧+间隔,空白段保留。""" # SUB 001 在 0/2s,4s 为空帧,SUB 002 在 6/8s。 frame_texts = [ (0.0, "SUB 001"), (2.0, "SUB 001"), (4.0, ""), (6.0, "SUB 002"), (8.0, "SUB 002"), ] frames = [] for index, (time, _text) in enumerate(frame_texts): image = tmp_path / f"f{index}.png" image.write_bytes(TEXT_IMG.read_bytes()) frames.append((time, image)) mapping = {str(image): text for (time, image), (_, text) in zip(frames, frame_texts)} def fake_vlm(node_id, request): return InvokeResponse(status="completed", outputs={"text": mapping[request.inputs["image_uri"]]}) monkeypatch.setattr("wov_app.registry.invoke", fake_vlm) manifest = _frames_manifest(tmp_path, frames) response = ocr_invoke( InvokeRequest( run_id="run_ocr", node_instance_id="", inputs={"frames_manifest": str(manifest)}, params={}, output_dir=str(tmp_path / "out"), ) ) assert response.status == "completed", response.error srt = Path(response.outputs["srt_uri"]).read_text(encoding="utf-8") assert srt.count("SUB 001") == 1 assert srt.count("SUB 002") == 1 # SUB 001 最后可见帧 2.0 + 间隔 2.0 = 4.0 消失(而非拖到 SUB 002 出现)。 assert "00:00:00,000 --> 00:00:04,000" in srt assert "00:00:06,000 --> 00:00:10,000" in srt def test_ocr_skips_failed_frames(monkeypatch, tmp_path) -> None: """个别帧 OCR 失败时跳过,不影响其余帧汇总。""" frames = [] for index in range(3): image = tmp_path / f"f{index}.png" image.write_bytes(TEXT_IMG.read_bytes()) frames.append((index * 2.0, image)) def fake_vlm(node_id, request): if "f1" in request.inputs["image_uri"]: return InvokeResponse(status="failed", error="boom") return InvokeResponse(status="completed", outputs={"text": "SUB 001"}) monkeypatch.setattr("wov_app.registry.invoke", fake_vlm) manifest = _frames_manifest(tmp_path, frames) response = ocr_invoke( InvokeRequest( run_id="run_ocr", node_instance_id="", inputs={"frames_manifest": str(manifest)}, params={}, output_dir=str(tmp_path / "out"), ) ) assert response.status == "completed", response.error srt = Path(response.outputs["srt_uri"]).read_text(encoding="utf-8") assert "SUB 001" in srt def test_ocr_missing_manifest(tmp_path) -> None: """缺少 frames_manifest 或清单文件不存在时返回失败。""" response = ocr_invoke( InvokeRequest(run_id="r", node_instance_id="", inputs={}, output_dir=str(tmp_path)) ) assert response.status == "failed" response = ocr_invoke( InvokeRequest( run_id="r", node_instance_id="", inputs={"frames_manifest": str(tmp_path / "none.json")}, output_dir=str(tmp_path), ) ) assert response.status == "failed" def test_ocr_skips_oversized_output(monkeypatch, tmp_path) -> None: """超长输出(模型重复循环等)直接报错跳过该帧,不进入 SRT。""" frames = [] for index in range(3): image = tmp_path / f"o{index}.png" image.write_bytes(TEXT_IMG.read_bytes()) frames.append((index * 2.0, image)) def fake_vlm(node_id, request): if "o0" in request.inputs["image_uri"]: return InvokeResponse(status="completed", outputs={"text": "重复字幕\n" * 50}) if "o1" in request.inputs["image_uri"]: # 空输出帧跳过。 return InvokeResponse(status="completed", outputs={"text": ""}) return InvokeResponse(status="completed", outputs={"text": "SUB 001"}) monkeypatch.setattr("wov_app.registry.invoke", fake_vlm) manifest = _frames_manifest(tmp_path, frames) response = ocr_invoke( InvokeRequest( run_id="run_ocr", node_instance_id="", inputs={"frames_manifest": str(manifest)}, params={"max_result_chars": 200}, output_dir=str(tmp_path / "out"), ) ) assert response.status == "completed", response.error srt = Path(response.outputs["srt_uri"]).read_text(encoding="utf-8") assert "SUB 001" in srt assert "重复字幕" not in srt def test_ocr_passes_short_text_through(monkeypatch, tmp_path) -> None: """不超过上限的模型输出原样进入 SRT(不再做应用层过滤)。""" image = tmp_path / "s0.png" image.write_bytes(TEXT_IMG.read_bytes()) frames = [(0.0, image)] def fake_vlm(node_id, request): return InvokeResponse(status="completed", outputs={"text": " SUB 001 "}) monkeypatch.setattr("wov_app.registry.invoke", fake_vlm) manifest = _frames_manifest(tmp_path, frames) response = ocr_invoke( InvokeRequest( run_id="run_ocr", node_instance_id="", inputs={"frames_manifest": str(manifest)}, params={}, output_dir=str(tmp_path / "out"), ) ) assert response.status == "completed", response.error srt = Path(response.outputs["srt_uri"]).read_text(encoding="utf-8") assert "SUB 001" in srt def test_ocr_interrupts_on_pause_flag(monkeypatch, tmp_path) -> None: """暂停信号:任务被暂停(paused.flag 存在)时 OCR 立即中断。 output_dir = /runs//steps/ocr,run 根目录为其父父目录; 调度器/API 在暂停时向 run 根写入 paused.flag,节点逐帧检查到即中止: 不调用 vlm-ocr、不写该帧断点存档,invoke 返回 failed(由调度器识别为 "被暂停"并保持 PAUSED,等待 resume 后从断点续跑)。 """ image = tmp_path / "s0.png" image.write_bytes(TEXT_IMG.read_bytes()) frames = [(i * 2.0, image) for i in range(4)] # 暂停信号位于 run 根目录(steps/ocr 的父父目录)。 run_root = tmp_path / "run_root" out_dir = run_root / "steps" / "ocr" out_dir.mkdir(parents=True) (run_root / "paused.flag").write_text("", encoding="utf-8") calls: list[str] = [] def fake_vlm(node_id, request): calls.append(request.inputs["image_uri"]) return InvokeResponse(status="completed", outputs={"text": "SUB 001"}) monkeypatch.setattr("wov_app.registry.invoke", fake_vlm) manifest = _frames_manifest(tmp_path, frames) response = ocr_invoke( InvokeRequest( run_id="run_paused_flag", node_instance_id="", inputs={"frames_manifest": str(manifest)}, params={}, output_dir=str(out_dir), ) ) # 节点以"被暂停"失败:调度器会据此保持 PAUSED 而不标 FAILED。 assert response.status == "failed" assert "暂停" in (response.error or "") assert calls == [] # 一帧都没有真正 OCR。 assert not (out_dir / "ocr_partial.jsonl").exists() # 未处理帧不入存档。 def test_format_eta() -> None: """预计剩余时间的格式化:秒/分/小时三种量级与边界值。""" from nodes.subtitle_ocr import _format_eta assert _format_eta(0) == "0秒" assert _format_eta(59.9) == "59秒" # 不足 1 分只显示秒。 assert _format_eta(61) == "1分01秒" assert _format_eta(2058.4) == "34分18秒" assert _format_eta(3600) == "1小时00分" # 1 小时整。 assert _format_eta(3725) == "1小时02分" # 超过 1 小时只显示到分钟。 assert _format_eta(-5) == "0秒" # 负数钳制为 0。 def test_eta_suffix() -> None: """进度日志的 ETA 后缀:速度为 0 不显示,速度正常时按剩余帧估算。""" from nodes.subtitle_ocr import _eta_suffix # 速率未知(0/负)→ 不显示 ETA。 assert _eta_suffix(done=100, total=100, rate=0) == "" assert _eta_suffix(done=100, total=100, rate=-1) == "" # 2673/22222 帧、9.5 帧/s:剩余 (22222-2673)/9.5 ≈ 2057.8s ≈ 34分。 assert _eta_suffix(done=2673, total=22222, rate=9.5) == ", 预计剩余 34分17秒" # 全部完成时剩余 0 秒。 assert _eta_suffix(done=22222, total=22222, rate=9.5) == ", 预计剩余 0秒"