edaa0d8290
- 从 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>
1000 lines
32 KiB
Python
1000 lines
32 KiB
Python
"""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
|