diff --git a/tests/integration/test_tree_build_e2e.py b/tests/integration/test_tree_build_e2e.py new file mode 100644 index 0000000..7885554 --- /dev/null +++ b/tests/integration/test_tree_build_e2e.py @@ -0,0 +1,172 @@ +"""建树模块端到端集成测试。 + +验证各模块协作: + 构造最小树 → verify → subtitle 注入 → TreeEnvironment 查询 → 序列化 roundtrip +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from app.tree.index import ( + IndexMeta, TreeIndex, L1Node, L1Card, + L2Node, L2Card, L3Node, L3Card, +) +from app.tree.verify import verify_tree +from app.tree.subtitle import SRTEntry, assign_subtitles_voronoi +from app.tree.environment import TreeEnvironment + + +class TestTreeModuleE2E: + def test_verify_subtitle_environment_pipeline(self, tmp_path): + """完整流程:构造树 → verify(删除幻觉实体)→ subtitle 注入 → environment 查询 → 序列化 roundtrip。""" + # 构造一棵树,L2 有混合实体(有出处/无出处) + l3_0 = L3Node( + id="vid_L1_000_L2_000_L3_000", + card=L3Card( + frame_summary="运动员在跑步冲刺", + visible_entities=["运动员", "跑道"], + ongoing_actions=["跑步"], + visible_text=["Nike", "2024"], + spatial_layout="居中构图", + visual_attributes={"lighting": "明亮", "camera_angle": "侧面"}, + ), + timestamp=2.0, + frame_path="frames/L1_000_L2_000_L3_000.jpg", + ) + l3_1 = L3Node( + id="vid_L1_000_L2_000_L3_001", + card=L3Card( + frame_summary="观众在看台上欢呼", + visible_entities=["观众", "看台"], + ongoing_actions=["欢呼"], + visible_text=["Stadium"], + spatial_layout="广角", + visual_attributes={}, + ), + timestamp=6.0, + frame_path="frames/L1_000_L2_000_L3_001.jpg", + ) + l2 = L2Node( + id="vid_L1_000_L2_000", + card=L2Card( + event_description="百米决赛片段", + entities=["运动员", "裁判", "幻觉实体XYZ"], + actions=["跑步", "欢呼"], + action_subjects=["运动员", "观众"], + visible_text=["Nike", "不存在的文字ABC"], + spatial_relations="运动员在跑道中央", + state_changes=None, + ), + time_range=(0.0, 10.0), + children=[l3_0, l3_1], + ) + l1 = L1Node( + id="vid_L1_000", + card=L1Card( + scene_summary="百米短跑决赛", + main_setting="体育场", + key_entities=["运动员", "不存在的人物"], + main_actions=["比赛"], + topic_keywords=["体育", "短跑"], + visible_text=["Nike", "Ghost文字"], + temporal_flow="从起跑到冲刺", + ), + time_range=(0.0, 10.0), + children=[l2], + ) + index = TreeIndex( + metadata=IndexMeta(source_path="/test/video.mp4", modality="video"), + roots=[l1], + ) + + # Step 1: verify — 删除无出处的实体和 visible_text + stats = verify_tree(index) + assert "幻觉实体XYZ" not in index.roots[0].children[0].card.entities + assert "不存在的文字ABC" not in index.roots[0].children[0].card.visible_text + assert "Ghost文字" not in index.roots[0].card.visible_text + assert "不存在的人物" not in index.roots[0].card.key_entities + # 有出处的保留 + assert "运动员" in index.roots[0].children[0].card.entities + assert "Nike" in index.roots[0].children[0].card.visible_text + + # Step 2: subtitle 注入 + srt_entries = [ + SRTEntry(start=1.0, end=3.0, text="And the runner sprints ahead!"), + SRTEntry(start=5.0, end=7.0, text="The crowd goes wild!"), + ] + assign_subtitles_voronoi(index, srt_entries) + assert l3_0.subtitle is not None + assert "sprints" in l3_0.subtitle + assert l3_1.subtitle is not None + assert "crowd" in l3_1.subtitle + + # Step 3: TreeEnvironment 查询 + env = TreeEnvironment(index) + + # view_node L3 + l3_view = env.view_node("vid_L1_000_L2_000_L3_000") + assert "运动员在跑步冲刺" in l3_view + + # view_node L2 (should list children) + l2_view = env.view_node("vid_L1_000_L2_000") + assert "百米决赛片段" in l2_view + assert "vid_L1_000_L2_000_L3_000" in l2_view + + # view_node with anchor + anchored = env.view_node("vid_L1_000_L2_000_L3_000", anchor=True) + assert "[c" in anchored + + # get_subtitle + assert "sprints" in env.get_subtitle("vid_L1_000_L2_000_L3_000") + + # search_similar (with embedding) + def fake_embed(texts): + if isinstance(texts, str): + texts = [texts] + rng = np.random.RandomState(42) + return rng.randn(len(texts), 4).astype(np.float32) + + index.embed_all(fake_embed, "test-model", 4) + results = env.search_similar("运动员跑步", top_k=3, embed_fn=fake_embed) + assert len(results) > 0 + + # Step 4: 序列化 roundtrip + path = tmp_path / "tree.json" + index.save_json(str(path)) + loaded = TreeIndex.load_json(str(path)) + + assert len(loaded.roots) == 1 + assert loaded.roots[0].card.scene_summary == "百米短跑决赛" + assert loaded.roots[0].children[0].children[0].subtitle is not None + assert "sprints" in loaded.roots[0].children[0].children[0].subtitle + # verify 的修改也被保留 + assert "幻觉实体XYZ" not in loaded.roots[0].children[0].card.entities + + def test_repair_detector_on_broken_tree(self): + """修复检测器能识别空卡片节点。""" + from app.tree.repair.detector import detect_issues + + l3 = L3Node( + id="vid_L1_000_L2_000_L3_000", + card=L3Card("", [], [], [], "", {}), # empty frame_summary + timestamp=1.0, + ) + l2 = L2Node( + id="vid_L1_000_L2_000", + card=L2Card("事件", [], [], [], [], "", None), + time_range=(0.0, 10.0), + children=[l3], + ) + l1 = L1Node( + id="vid_L1_000", + card=L1Card("场景", "", [], [], [], [], ""), + time_range=(0.0, 10.0), + children=[l2], + ) + index = TreeIndex(metadata=IndexMeta("/t.mp4", "video"), roots=[l1]) + + issues = detect_issues(index) + assert len(issues) >= 1 + assert any(i.issue_type == "empty_field" for i in issues)