- 从 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>
This commit is contained in:
@@ -0,0 +1,999 @@
|
||||
"""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 帧/clip（5 秒 * 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 返回无效 JSON，single 返回有效 JSON，L2 正常。"""
|
||||
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
|
||||
Reference in New Issue
Block a user