Files
Video-Tree-TRM5/tests/unit/test_video_builder.py
T
iomgaa edaa0d8290 feat(tree): VideoTreeBuilder 保真 #1 #2 #3 + 复杂度重构
- 从 reference/video_tree_trm/video_tree_builder.py (994行) 迁移
- 保真算法 #1: L2 轴心建树策略 (asyncio.gather 链式并发)
- 保真算法 #2: VLM 批量帧描述 + JSON fallback (_L3_BATCH_SIZE=5)
- 保真算法 #3: 断点续跑 (progress.json + L1 中间 JSON)
- 新增: VLMProvider/LLMProvider Protocol 替代 LLMClient
- 新增: 结构化 JSON 输出 → L1Card/L2Card/L3Card
- 新增: L2 代表帧复用 L3 帧 (_sample_representative_frames)
- 新增: 字幕注入 + Voronoi 分配
- 重构: 提取 _load_resume_state/_assemble_roots 降低 _build_async 复杂度 D(21)→C(14)
- 44 个单元测试全部通过

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-07-07 02:17:50 -04:00

1000 lines
32 KiB
Python
Raw 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.
"""VideoTreeBuilder 单元测试。
测试覆盖:
- _segment_video: 时间切分
- _get_l2_clips: 片段切分
- _sample_representative_frames: 帧采样
- _parse_l3_cards: L3 JSON 解析 + fallback 条件
- _parse_l2_card: L2 JSON 解析
- _parse_l1_card: L1 JSON 解析
- _parse_l3_card_single: 单帧 L3 JSON 解析
- build: 完整构建流程(mock VLM/LLM
- checkpoint/resume: 断点续跑机制
"""
from __future__ import annotations
import json
from pathlib import Path
from typing import Any
from unittest.mock import patch
import pytest
from app.tree.config import TreeConfig
from app.tree.index import (
L1Card,
L1Node,
L2Card,
L2Node,
L3Card,
L3Node,
)
from app.tree.video_builder import VideoTreeBuilder
from core.types import LLMResponse
# ---------------------------------------------------------------------------
# Mock 提供者
# ---------------------------------------------------------------------------
def _make_llm_response(content: str) -> LLMResponse:
"""构造 LLMResponse 测试工具。"""
return LLMResponse(
content=content,
thinking="",
model="mock",
provider="mock",
prompt_tokens=10,
completion_tokens=10,
latency_ms=10,
ttft_ms=1.0,
max_inter_token_ms=1.0,
cache_hit=False,
call_id="mock-call",
)
def _l3_card_dict(idx: int = 0) -> dict[str, Any]:
"""构造单个 L3Card 的字典表示。"""
return {
"frame_summary": f"帧 {idx} 的描述",
"visible_entities": ["实体A"],
"ongoing_actions": ["动作A"],
"visible_text": [],
"spatial_layout": "居中",
"visual_attributes": {
"lighting": "自然光",
"dominant_colors": ["白"],
"camera_angle": "正面",
},
}
def _l2_card_dict() -> dict[str, Any]:
"""构造 L2Card 的字典表示。"""
return {
"event_description": "视频片段描述",
"entities": ["实体A"],
"actions": ["动作A"],
"action_subjects": ["主体A"],
"visible_text": [],
"spatial_relations": "居中",
"state_changes": None,
}
def _l1_card_dict() -> dict[str, Any]:
"""构造 L1Card 的字典表示。"""
return {
"scene_summary": "场景摘要描述",
"main_setting": "室内",
"key_entities": ["实体A"],
"main_actions": ["动作A"],
"topic_keywords": ["关键词A"],
"visible_text": [],
"temporal_flow": "从开始到结束",
}
class MockVLMProvider:
"""模拟 VLM 提供者,根据 prompt 类型返回固定 JSON 响应。"""
def __init__(self) -> None:
self.calls: list[dict[str, Any]] = []
async def chat_with_images(
self,
messages: list[dict[str, Any]],
images: list[str | Path],
*,
session_id: str | None = None,
parent_call_id: str | None = None,
) -> LLMResponse:
"""根据 prompt 内容判断调用类型,返回对应 JSON。"""
self.calls.append({"messages": messages, "images": images})
content = messages[0]["content"]
n_images = len(images)
if "JSON 数组" in content:
# L3 batch prompt
cards = [_l3_card_dict(i) for i in range(n_images)]
return _make_llm_response(json.dumps(cards, ensure_ascii=False))
if "用一到两句话描述这帧" in content:
# L3 single prompt
return _make_llm_response(
json.dumps(_l3_card_dict(0), ensure_ascii=False),
)
# L2 prompt
return _make_llm_response(
json.dumps(_l2_card_dict(), ensure_ascii=False),
)
class MockLLMProvider:
"""模拟 LLM 提供者,返回固定 L1Card JSON 响应。"""
def __init__(self) -> None:
self.calls: list[dict[str, Any]] = []
async def chat(
self,
messages: list[dict[str, Any]],
*,
session_id: str | None = None,
parent_call_id: str | None = None,
) -> LLMResponse:
"""返回 L1Card JSON。"""
self.calls.append({"messages": messages})
return _make_llm_response(
json.dumps(_l1_card_dict(), ensure_ascii=False),
)
# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------
@pytest.fixture()
def tree_config(tmp_path: Path) -> TreeConfig:
"""简化配置:10 秒视频→1 个 L1 段→2 个 L2 clip→每 clip 5 帧。"""
return TreeConfig(
l1_segment_duration=10.0,
l2_clip_duration=5.0,
l3_fps=1.0,
l2_representative_frames=2,
cache_dir=str(tmp_path / "cache"),
concurrency=4,
)
@pytest.fixture()
def mock_vlm() -> MockVLMProvider:
"""返回 MockVLMProvider 实例。"""
return MockVLMProvider()
@pytest.fixture()
def mock_llm() -> MockLLMProvider:
"""返回 MockLLMProvider 实例。"""
return MockLLMProvider()
@pytest.fixture()
def builder(
mock_vlm: MockVLMProvider,
mock_llm: MockLLMProvider,
tree_config: TreeConfig,
) -> VideoTreeBuilder:
"""构造带 mock 依赖的 VideoTreeBuilder。"""
return VideoTreeBuilder(vlm=mock_vlm, llm=mock_llm, config=tree_config)
# ---------------------------------------------------------------------------
# 测试:_segment_video
# ---------------------------------------------------------------------------
class TestSegmentVideo:
"""测试时间切分逻辑。"""
def test_exact_division(self, builder: VideoTreeBuilder) -> None:
"""总时长能被 L1 段时长整除时的切分。"""
ranges = builder._segment_video("dummy", duration_hint=20.0)
assert ranges == [(0.0, 10.0), (10.0, 20.0)]
def test_non_divisible(self, builder: VideoTreeBuilder) -> None:
"""总时长不能被整除时,末段截断。"""
ranges = builder._segment_video("dummy", duration_hint=15.0)
assert len(ranges) == 2
assert ranges[0] == (0.0, 10.0)
assert ranges[1] == (10.0, 15.0)
def test_short_video(self, builder: VideoTreeBuilder) -> None:
"""短视频(时长 < L1 段时长)产生单段。"""
ranges = builder._segment_video("dummy", duration_hint=5.0)
assert ranges == [(0.0, 5.0)]
# ---------------------------------------------------------------------------
# 测试:_get_l2_clips
# ---------------------------------------------------------------------------
class TestGetL2Clips:
"""测试 L2 clip 切分逻辑。"""
def test_basic_clips(self, builder: VideoTreeBuilder) -> None:
"""L1 区间能被 L2 步长整除。"""
clips = builder._get_l2_clips((0.0, 10.0))
assert clips == [(0.0, 5.0), (5.0, 10.0)]
def test_remainder_clip(self, builder: VideoTreeBuilder) -> None:
"""L1 区间不能被整除时,末段截断。"""
clips = builder._get_l2_clips((0.0, 7.0))
assert len(clips) == 2
assert clips[0] == (0.0, 5.0)
assert clips[1] == (5.0, 7.0)
# ---------------------------------------------------------------------------
# 测试:_sample_representative_frames
# ---------------------------------------------------------------------------
class TestSampleRepresentativeFrames:
"""测试帧采样逻辑。"""
def test_fewer_frames_than_n(self) -> None:
"""帧数不足时返回全部。"""
frames = [("a.jpg", 0.0), ("b.jpg", 1.0)]
result = VideoTreeBuilder._sample_representative_frames(frames, 5)
assert result == ["a.jpg", "b.jpg"]
def test_exact_n(self) -> None:
"""帧数恰等于 n 时返回全部。"""
frames = [("a.jpg", 0.0), ("b.jpg", 1.0), ("c.jpg", 2.0)]
result = VideoTreeBuilder._sample_representative_frames(frames, 3)
assert result == ["a.jpg", "b.jpg", "c.jpg"]
def test_uniform_sampling(self) -> None:
"""10 帧中采样 3 帧,应均匀分布。"""
frames = [(f"{i}.jpg", float(i)) for i in range(10)]
result = VideoTreeBuilder._sample_representative_frames(frames, 3)
# step = 10/3 = 3.33 → indices 0, 3, 6
assert result == ["0.jpg", "3.jpg", "6.jpg"]
def test_sampling_two_from_five(self) -> None:
"""5 帧中采样 2 帧。"""
frames = [(f"{i}.jpg", float(i)) for i in range(5)]
result = VideoTreeBuilder._sample_representative_frames(frames, 2)
# step = 5/2 = 2.5 → indices 0, 2
assert result == ["0.jpg", "2.jpg"]
# ---------------------------------------------------------------------------
# 测试:_parse_l3_cards
# ---------------------------------------------------------------------------
class TestParseL3Cards:
"""测试 L3 批量 JSON 解析(保真算法 #2 的解析环节)。"""
def test_valid_json_array(self, builder: VideoTreeBuilder) -> None:
"""正常 JSON 数组解析成功。"""
raw = json.dumps([_l3_card_dict(i) for i in range(3)])
result = builder._parse_l3_cards(raw, 3)
assert result is not None
assert len(result) == 3
assert result[0].frame_summary == "帧 0 的描述"
assert result[2].frame_summary == "帧 2 的描述"
def test_count_mismatch_returns_none(
self,
builder: VideoTreeBuilder,
) -> None:
"""数量不匹配 → 返回 None(触发 fallback)。"""
raw = json.dumps([_l3_card_dict(0), _l3_card_dict(1)])
result = builder._parse_l3_cards(raw, 3)
assert result is None
def test_missing_field_returns_none(
self,
builder: VideoTreeBuilder,
) -> None:
"""必填字段缺失 → 整批次返回 None。"""
card = _l3_card_dict(0)
del card["frame_summary"]
raw = json.dumps([card])
result = builder._parse_l3_cards(raw, 1)
assert result is None
def test_invalid_json_returns_none(
self,
builder: VideoTreeBuilder,
) -> None:
"""非法 JSON 返回 None。"""
result = builder._parse_l3_cards("not json at all", 1)
assert result is None
def test_json_in_code_block(self, builder: VideoTreeBuilder) -> None:
"""Markdown 代码块包裹的 JSON 也能解析。"""
inner = json.dumps([_l3_card_dict(0)])
raw = f"```json\n{inner}\n```"
result = builder._parse_l3_cards(raw, 1)
assert result is not None
assert len(result) == 1
def test_non_dict_item_returns_none(
self,
builder: VideoTreeBuilder,
) -> None:
"""数组元素非 dict → 返回 None。"""
raw = json.dumps(["string_item"])
result = builder._parse_l3_cards(raw, 1)
assert result is None
# ---------------------------------------------------------------------------
# 测试:_parse_l2_card
# ---------------------------------------------------------------------------
class TestParseL2Card:
"""测试 L2 JSON 解析。"""
def test_valid_json(self, builder: VideoTreeBuilder) -> None:
"""正常 JSON 解析成功。"""
raw = json.dumps(_l2_card_dict(), ensure_ascii=False)
card = builder._parse_l2_card(raw)
assert card.event_description == "视频片段描述"
assert card.entities == ["实体A"]
assert card.state_changes is None
def test_invalid_json_fallback(
self,
builder: VideoTreeBuilder,
) -> None:
"""JSON 解析失败 → 退化卡片。"""
card = builder._parse_l2_card("这是一段普通文字描述")
assert card.event_description == "这是一段普通文字描述"
assert card.entities == []
def test_with_state_changes(self, builder: VideoTreeBuilder) -> None:
"""state_changes 非 null 时正常解析。"""
d = _l2_card_dict()
d["state_changes"] = "从站立到坐下"
raw = json.dumps(d, ensure_ascii=False)
card = builder._parse_l2_card(raw)
assert card.state_changes == "从站立到坐下"
# ---------------------------------------------------------------------------
# 测试:_parse_l1_card
# ---------------------------------------------------------------------------
class TestParseL1Card:
"""测试 L1 JSON 解析。"""
def test_valid_json(self, builder: VideoTreeBuilder) -> None:
"""正常 JSON 解析成功。"""
raw = json.dumps(_l1_card_dict(), ensure_ascii=False)
card = builder._parse_l1_card(raw)
assert card.scene_summary == "场景摘要描述"
assert card.main_setting == "室内"
def test_invalid_json_fallback(
self,
builder: VideoTreeBuilder,
) -> None:
"""JSON 解析失败 → 退化卡片。"""
card = builder._parse_l1_card("这是段落摘要")
assert card.scene_summary == "这是段落摘要"
assert card.key_entities == []
# ---------------------------------------------------------------------------
# 测试:_parse_l3_card_single
# ---------------------------------------------------------------------------
class TestParseL3CardSingle:
"""测试单帧 L3 JSON 解析。"""
def test_valid_json(self, builder: VideoTreeBuilder) -> None:
"""正常 JSON 解析成功。"""
raw = json.dumps(_l3_card_dict(0), ensure_ascii=False)
card = builder._parse_l3_card_single(raw)
assert card.frame_summary == "帧 0 的描述"
def test_invalid_json_fallback(
self,
builder: VideoTreeBuilder,
) -> None:
"""JSON 解析失败 → 退化卡片。"""
card = builder._parse_l3_card_single("一帧画面的描述")
assert card.frame_summary == "一帧画面的描述"
assert card.visible_entities == []
# ---------------------------------------------------------------------------
# 测试:_extract_json
# ---------------------------------------------------------------------------
class TestExtractJson:
"""测试 JSON 提取辅助方法。"""
def test_plain_object(self) -> None:
"""直接 JSON 对象。"""
result = VideoTreeBuilder._extract_json('{"key": "value"}')
assert result == {"key": "value"}
def test_plain_array(self) -> None:
"""直接 JSON 数组。"""
result = VideoTreeBuilder._extract_json("[1, 2, 3]")
assert result == [1, 2, 3]
def test_code_block(self) -> None:
"""Markdown 代码块中的 JSON。"""
raw = '```json\n{"key": "value"}\n```'
result = VideoTreeBuilder._extract_json(raw)
assert result == {"key": "value"}
def test_surrounding_text(self) -> None:
"""JSON 前后有文字。"""
raw = 'Here is the result: {"key": "value"} end'
result = VideoTreeBuilder._extract_json(raw)
assert result == {"key": "value"}
def test_invalid(self) -> None:
"""完全无 JSON 内容。"""
result = VideoTreeBuilder._extract_json("no json here")
assert result is None
# ---------------------------------------------------------------------------
# 测试:完整构建流程
# ---------------------------------------------------------------------------
def _mock_ffmpeg_factory(tmp_path: Path):
"""创建 mock ffmpeg 帧提取函数。"""
def _mock_extract(
video_path: str,
ts: float,
out_path: str,
) -> bool:
Path(out_path).parent.mkdir(parents=True, exist_ok=True)
Path(out_path).write_bytes(b"FAKE_JPEG")
return True
return _mock_extract
class TestBuildFullFlow:
"""测试完整构建流程(mock VLM/LLM/ffmpeg)。"""
def test_build_produces_correct_structure(
self,
mock_vlm: MockVLMProvider,
mock_llm: MockLLMProvider,
tmp_path: Path,
) -> None:
"""10 秒视频 → 1 L1 → 2 L2 → 每 L2 约 5 帧 L3。"""
config = TreeConfig(
l1_segment_duration=10.0,
l2_clip_duration=5.0,
l3_fps=1.0,
l2_representative_frames=2,
cache_dir=str(tmp_path / "cache"),
concurrency=4,
)
builder = VideoTreeBuilder(
vlm=mock_vlm,
llm=mock_llm,
config=config,
)
dummy_video = tmp_path / "test_video.mp4"
dummy_video.write_bytes(b"FAKE")
with (
patch.object(
builder,
"_segment_video",
return_value=[(0.0, 10.0)],
),
patch.object(
builder,
"_ffmpeg_extract_frame",
side_effect=_mock_ffmpeg_factory(tmp_path),
),
):
index = builder.build(str(dummy_video))
# 结构校验
assert len(index.roots) == 1
l1 = index.roots[0]
assert l1.id == "l1_0"
assert l1.card.scene_summary == "场景摘要描述"
assert l1.time_range == (0.0, 10.0)
# 2 个 L2 clip
assert len(l1.children) == 2
for j, l2 in enumerate(l1.children):
assert l2.id == f"l1_0_l2_{j}"
assert l2.card.event_description == "视频片段描述"
# 5 帧/clip5 秒 * 1 fps
assert len(l2.children) == 5
for k, l3 in enumerate(l2.children):
assert l3.id == f"l1_0_l2_{j}_l3_{k}"
assert l3.card.frame_summary is not None
assert l3.frame_path is not None
assert l3.timestamp is not None
# 确认 VLM/LLM 被调用
# L2: 2 calls (one per clip)
# L3: 2 calls (one batch per clip, each batch has 5 frames)
# L1: 1 call
assert len(mock_vlm.calls) == 4 # 2 L2 + 2 L3 batches
assert len(mock_llm.calls) == 1 # 1 L1
# metadata 校验
assert index.metadata.source_path == str(dummy_video)
assert index.metadata.modality == "video"
def test_build_cleans_up_intermediate(
self,
mock_vlm: MockVLMProvider,
mock_llm: MockLLMProvider,
tmp_path: Path,
) -> None:
"""构建成功后中间文件已清理。"""
config = TreeConfig(
l1_segment_duration=10.0,
l2_clip_duration=10.0,
l3_fps=1.0,
l2_representative_frames=2,
cache_dir=str(tmp_path / "cache"),
concurrency=4,
)
builder = VideoTreeBuilder(
vlm=mock_vlm,
llm=mock_llm,
config=config,
)
dummy_video = tmp_path / "test_video.mp4"
dummy_video.write_bytes(b"FAKE")
with (
patch.object(
builder,
"_segment_video",
return_value=[(0.0, 10.0)],
),
patch.object(
builder,
"_ffmpeg_extract_frame",
side_effect=_mock_ffmpeg_factory(tmp_path),
),
):
builder.build(str(dummy_video))
# 中间文件应已清理
progress_dir = tmp_path / "cache" / "progress"
inter_dir = tmp_path / "cache" / "intermediate" / "test_video"
assert not (progress_dir / "test_video.json").exists()
# intermediate 目录可能不存在或为空
if inter_dir.exists():
assert len(list(inter_dir.glob("l1_*.json"))) == 0
# ---------------------------------------------------------------------------
# 测试:L3 fallback(保真算法 #2
# ---------------------------------------------------------------------------
class MockVLMWithBatchFailure:
"""模拟批量 VLM 调用失败、单帧调用成功的 VLM 提供者。"""
def __init__(self) -> None:
self.calls: list[dict[str, Any]] = []
async def chat_with_images(
self,
messages: list[dict[str, Any]],
images: list[str | Path],
*,
session_id: str | None = None,
parent_call_id: str | None = None,
) -> LLMResponse:
"""batch 返回无效 JSONsingle 返回有效 JSONL2 正常。"""
self.calls.append({"n_images": len(images)})
content = messages[0]["content"]
if "JSON 数组" in content:
# L3 batch: 返回无效 JSON 触发 fallback
return _make_llm_response("INVALID JSON OUTPUT")
if "用一到两句话描述这帧" in content:
# L3 single fallback: 有效 JSON
return _make_llm_response(
json.dumps(_l3_card_dict(0), ensure_ascii=False),
)
# L2: 有效 JSON
return _make_llm_response(
json.dumps(_l2_card_dict(), ensure_ascii=False),
)
class TestL3Fallback:
"""测试 L3 批量失败→逐帧 fallback(保真算法 #2)。"""
def test_fallback_to_single_frame(
self,
mock_llm: MockLLMProvider,
tmp_path: Path,
) -> None:
"""批量 VLM 解析失败时,逐帧 fallback 仍能构建完整树。"""
vlm = MockVLMWithBatchFailure()
config = TreeConfig(
l1_segment_duration=10.0,
l2_clip_duration=10.0,
l3_fps=1.0,
l2_representative_frames=2,
cache_dir=str(tmp_path / "cache"),
concurrency=4,
)
builder = VideoTreeBuilder(vlm=vlm, llm=mock_llm, config=config)
dummy_video = tmp_path / "test_video.mp4"
dummy_video.write_bytes(b"FAKE")
with (
patch.object(
builder,
"_segment_video",
return_value=[(0.0, 10.0)],
),
patch.object(
builder,
"_ffmpeg_extract_frame",
side_effect=_mock_ffmpeg_factory(tmp_path),
),
):
index = builder.build(str(dummy_video))
# 结构仍然完整
assert len(index.roots) == 1
l1 = index.roots[0]
assert len(l1.children) == 1 # 1 clip (10s clip)
l2 = l1.children[0]
# 10 frames (10s * 1fps), all from single-frame fallback
assert len(l2.children) == 10
for l3 in l2.children:
assert l3.card.frame_summary == "帧 0 的描述"
# VLM 调用次数:1 L2 + 1 batch(fail) + 10 single = 12
# 但 batch 分为 2 batches (10 frames / 5 per batch)
# 所以: 1 L2 + 2 batch(fail) + 10 single = 13
assert len(vlm.calls) == 13
# ---------------------------------------------------------------------------
# 测试:断点续跑(保真算法 #3
# ---------------------------------------------------------------------------
class TestCheckpointResume:
"""测试断点续跑机制。"""
def test_save_and_load_progress(
self,
builder: VideoTreeBuilder,
) -> None:
"""进度文件保存和加载。"""
stem = "test_video"
builder._save_progress(stem, total_l1=3, finished_l1_ids={0, 1})
progress = builder._load_progress(stem)
assert progress is not None
assert progress["total_l1"] == 3
assert sorted(progress["finished_l1_ids"]) == [0, 1]
assert "created_at" in progress
assert "updated_at" in progress
def test_load_nonexistent_progress(
self,
builder: VideoTreeBuilder,
) -> None:
"""不存在的进度文件返回 None。"""
assert builder._load_progress("nonexistent") is None
def test_save_and_load_l1_intermediate(
self,
builder: VideoTreeBuilder,
) -> None:
"""L1 中间结果保存和加载。"""
stem = "test_video"
l1_card = L1Card(
scene_summary="测试摘要",
main_setting="测试场景",
key_entities=["实体"],
main_actions=["动作"],
topic_keywords=["关键词"],
visible_text=[],
temporal_flow="测试流向",
)
l2_card = L2Card(
event_description="事件描述",
entities=["实体"],
actions=["动作"],
action_subjects=["主体"],
visible_text=[],
spatial_relations="居中",
state_changes=None,
)
l3_card = L3Card(
frame_summary="帧描述",
visible_entities=["实体"],
ongoing_actions=["动作"],
visible_text=[],
spatial_layout="居中",
visual_attributes={"lighting": "自然光"},
)
l3_node = L3Node(
id="l1_0_l2_0_l3_0",
card=l3_card,
timestamp=1.0,
frame_path="/tmp/frame.jpg",
)
l2_node = L2Node(
id="l1_0_l2_0",
card=l2_card,
time_range=(0.0, 5.0),
children=[l3_node],
)
l1_node = L1Node(
id="l1_0",
card=l1_card,
time_range=(0.0, 10.0),
children=[l2_node],
)
builder._save_l1_intermediate(stem, l1_node, 0)
assert builder._has_l1_intermediate(stem, 0)
assert not builder._has_l1_intermediate(stem, 1)
loaded = builder._load_l1_intermediate(stem, 0)
assert loaded is not None
assert loaded.id == "l1_0"
assert loaded.card.scene_summary == "测试摘要"
assert len(loaded.children) == 1
assert loaded.children[0].id == "l1_0_l2_0"
def test_cleanup_removes_files(
self,
builder: VideoTreeBuilder,
) -> None:
"""清理函数删除进度文件和中间 JSON。"""
stem = "test_video"
builder._save_progress(stem, total_l1=1, finished_l1_ids={0})
# 创建一个假的中间文件
inter_dir = builder._intermediate_dir(stem)
inter_dir.mkdir(parents=True, exist_ok=True)
(inter_dir / "l1_0.json").write_text("{}")
builder._cleanup_intermediate_and_progress(stem)
assert not builder._progress_path(stem).is_file()
assert not (inter_dir / "l1_0.json").is_file()
def test_resume_skips_finished_segments(
self,
mock_vlm: MockVLMProvider,
mock_llm: MockLLMProvider,
tmp_path: Path,
) -> None:
"""断点续跑:跳过已完成的 L1 段,只构建未完成的段。"""
config = TreeConfig(
l1_segment_duration=5.0,
l2_clip_duration=5.0,
l3_fps=1.0,
l2_representative_frames=2,
cache_dir=str(tmp_path / "cache"),
concurrency=4,
)
builder = VideoTreeBuilder(
vlm=mock_vlm,
llm=mock_llm,
config=config,
)
source_id = "test_resume"
# Phase 1: 手动创建 L1_0 的中间结果(模拟已完成)
l1_card = L1Card(
scene_summary="已完成的段",
main_setting="场景A",
key_entities=["实体A"],
main_actions=["动作A"],
topic_keywords=["关键词A"],
visible_text=[],
temporal_flow="流向A",
)
l2_card = L2Card(
event_description="已完成的片段",
entities=["实体A"],
actions=["动作A"],
action_subjects=["主体A"],
visible_text=[],
spatial_relations="居中",
state_changes=None,
)
l3_card = L3Card(
frame_summary="已完成的帧",
visible_entities=["实体A"],
ongoing_actions=["动作A"],
visible_text=[],
spatial_layout="居中",
visual_attributes={"lighting": "自然光"},
)
l3_node = L3Node(
id="l1_0_l2_0_l3_0",
card=l3_card,
timestamp=1.0,
)
l2_node = L2Node(
id="l1_0_l2_0",
card=l2_card,
time_range=(0.0, 5.0),
children=[l3_node],
)
l1_node_0 = L1Node(
id="l1_0",
card=l1_card,
time_range=(0.0, 5.0),
children=[l2_node],
)
builder._save_l1_intermediate(source_id, l1_node_0, 0)
# Phase 2: 创建进度文件(标记 L1_0 完成)
builder._save_progress(
source_id,
total_l1=2,
finished_l1_ids={0},
)
# Phase 3: 构建(应跳过 L1_0,只构建 L1_1)
dummy_video = tmp_path / f"{source_id}.mp4"
dummy_video.write_bytes(b"FAKE")
with (
patch.object(
builder,
"_segment_video",
return_value=[(0.0, 5.0), (5.0, 10.0)],
),
patch.object(
builder,
"_ffmpeg_extract_frame",
side_effect=_mock_ffmpeg_factory(tmp_path),
),
):
index = builder.build(str(dummy_video))
# Phase 4: 验证
assert len(index.roots) == 2
# L1_0 来自中间结果
assert index.roots[0].id == "l1_0"
assert index.roots[0].card.scene_summary == "已完成的段"
# L1_1 是新构建的
assert index.roots[1].id == "l1_1"
assert index.roots[1].card.scene_summary == "场景摘要描述"
# VLM 只被调用了 L1_1 的部分(1 L2 + 1 L3 batch
assert len(mock_vlm.calls) == 2 # 1 L2 + 1 L3
assert len(mock_llm.calls) == 1 # 1 L1
# 构建完成后,进度和中间文件已清理
assert not builder._progress_path(source_id).is_file()
# ---------------------------------------------------------------------------
# 测试:字幕注入
# ---------------------------------------------------------------------------
class TestSubtitleInjection:
"""测试字幕注入功能。"""
def test_build_subtitle_block_with_entries(
self,
builder: VideoTreeBuilder,
) -> None:
"""有匹配字幕时返回字幕文本块。"""
from app.tree.subtitle import SRTEntry
entries = [
SRTEntry(start=1.0, end=3.0, text="你好世界"),
SRTEntry(start=4.0, end=6.0, text="再见"),
]
block = builder._build_subtitle_block(entries, (0.0, 5.0))
assert "字幕信息" in block
assert "你好世界" in block
def test_build_subtitle_block_no_match(
self,
builder: VideoTreeBuilder,
) -> None:
"""无匹配字幕时返回空字符串。"""
from app.tree.subtitle import SRTEntry
entries = [SRTEntry(start=100.0, end=110.0, text="远处的字幕")]
block = builder._build_subtitle_block(entries, (0.0, 5.0))
assert block == ""
def test_build_subtitle_block_none_entries(
self,
builder: VideoTreeBuilder,
) -> None:
"""srt_entries 为 None 时返回空字符串。"""
block = builder._build_subtitle_block(None, (0.0, 5.0))
assert block == ""
def test_build_subtitle_block_point_range(
self,
builder: VideoTreeBuilder,
) -> None:
"""点时间范围(单帧)自动扩展窗口。"""
from app.tree.subtitle import SRTEntry
entries = [SRTEntry(start=4.0, end=6.0, text="窗口内字幕")]
# 点时间 5.0,窗口 ±5.0 → (0.0, 10.0)
block = builder._build_subtitle_block(entries, (5.0, 5.0))
assert "窗口内字幕" in block
# ---------------------------------------------------------------------------
# 测试:URL 与 stem 辅助
# ---------------------------------------------------------------------------
class TestHelpers:
"""测试静态辅助方法。"""
def test_is_url_true(self) -> None:
"""HTTP/HTTPS URL 识别。"""
assert VideoTreeBuilder._is_url("https://example.com/video.mp4")
assert VideoTreeBuilder._is_url("http://example.com/video.mp4")
def test_is_url_false(self) -> None:
"""本地路径不是 URL。"""
assert not VideoTreeBuilder._is_url("/path/to/video.mp4")
assert not VideoTreeBuilder._is_url("video.mp4")
def test_source_stem_local(self) -> None:
"""本地文件的 stem。"""
assert VideoTreeBuilder._source_stem("/path/to/my_video.mp4") == "my_video"
def test_source_stem_youtube(self) -> None:
"""YouTube URL 的 stem 为视频 ID。"""
stem = VideoTreeBuilder._source_stem(
"https://www.youtube.com/watch?v=dQw4w9WgXcQ",
)
assert stem == "dQw4w9WgXcQ"
def test_source_stem_long_name(self) -> None:
"""超长文件名截断到 64 字符。"""
long_name = "a" * 100 + ".mp4"
stem = VideoTreeBuilder._source_stem(f"/path/{long_name}")
assert len(stem) == 64