"""TreeEnvironment 运行时单元测试。""" from __future__ import annotations from pathlib import Path import numpy as np import pytest from app.tree.environment import TreeEnvironment from app.tree.index import ( IndexMeta, L1Card, L1Node, L2Card, L2Node, L3Card, L3Node, TreeIndex, ) def _make_test_index() -> TreeIndex: """构建测试用的三层树索引。""" l3_0 = L3Node( id="vid_L1_000_L2_000_L3_000", card=L3Card( "运动员在跑步", ["运动员"], ["跑步"], ["Nike"], "居中", {"lighting": "明亮"}, ), timestamp=1.0, frame_path="frames/L1_000_L2_000_L3_000.jpg", subtitle="he is running", ) l3_1 = L3Node( id="vid_L1_000_L2_000_L3_001", card=L3Card( "观众欢呼", ["观众"], ["欢呼"], [], "广角", {}, ), timestamp=3.0, frame_path="frames/L1_000_L2_000_L3_001.jpg", ) l2 = L2Node( id="vid_L1_000_L2_000", card=L2Card( "比赛片段", ["运动员"], ["跑步"], ["运动员"], ["Nike"], "", None, ), time_range=(0.0, 10.0), children=[l3_0, l3_1], ) l1 = L1Node( id="vid_L1_000", card=L1Card( "体育赛事", "体育场", ["运动员"], ["比赛"], ["体育"], ["Nike"], "从左到右", ), time_range=(0.0, 10.0), children=[l2], ) return TreeIndex(metadata=IndexMeta("/test.mp4", "video"), roots=[l1]) class TestViewNode: """view_node 方法测试。""" def test_l3_node(self) -> None: """L3 节点应显示帧描述。""" env = TreeEnvironment(_make_test_index()) result = env.view_node("vid_L1_000_L2_000_L3_000") assert "运动员在跑步" in result assert "vid_L1_000_L2_000_L3_000" in result def test_l2_node_shows_children(self) -> None: """L2 节点应显示子节点概览。""" env = TreeEnvironment(_make_test_index()) result = env.view_node("vid_L1_000_L2_000") assert "比赛片段" in result assert "vid_L1_000_L2_000_L3_000" in result # child listed def test_anchor_mode(self) -> None: """锚模式应在输出中添加 [cN] 标记。""" env = TreeEnvironment(_make_test_index()) result = env.view_node("vid_L1_000_L2_000_L3_000", anchor=True) assert "[c" in result # anchor markers present def test_unknown_node_raises(self) -> None: """查询不存在的节点应抛出 KeyError。""" env = TreeEnvironment(_make_test_index()) with pytest.raises(KeyError): env.view_node("nonexistent") class TestSearchSimilar: """search_similar 方法测试。""" def test_returns_results(self) -> None: """使用 embed_fn 应返回搜索结果。""" index = _make_test_index() def fake_embed( texts: str | list[str], ) -> np.ndarray: if isinstance(texts, str): texts = [texts] return np.random.randn(len(texts), 4).astype(np.float32) index.embed_all(fake_embed, "test", 4) env = TreeEnvironment(index) results = env.search_similar("运动员", top_k=3, embed_fn=fake_embed) assert len(results) > 0 assert all(isinstance(r, tuple) and len(r) == 2 for r in results) def test_ancestor_dedup(self) -> None: """祖先去重:如果 L3 已在结果中,其 L1/L2 祖先应被跳过。""" index = _make_test_index() # 手动设置 embedding,使 L3 节点分数高于 L1/L2 l3_0 = index.roots[0].children[0].children[0] l3_1 = index.roots[0].children[0].children[1] l2 = index.roots[0].children[0] l1 = index.roots[0] l3_0.embedding = np.array([1.0, 0, 0, 0], dtype=np.float32) l3_1.embedding = np.array([0.9, 0.1, 0, 0], dtype=np.float32) l2.embedding = np.array([0.5, 0.5, 0, 0], dtype=np.float32) l1.embedding = np.array([0.3, 0.3, 0.3, 0], dtype=np.float32) index.metadata.embed_model = "test" index.metadata.embed_dim = 4 env = TreeEnvironment(index) results = env.search_similar( "运动员", top_k=5, embed_fn=lambda t: np.array([[1.0, 0, 0, 0]], dtype=np.float32), ) result_ids = [r[0] for r in results] # L3 节点应存在;其祖先应被去重跳过 assert "vid_L1_000_L2_000_L3_000" in result_ids def test_with_embed_fn_overrides_existing(self) -> None: """即使已有 embedding,提供 embed_fn 时仍应用于 query 编码。""" index = _make_test_index() def fake_embed( texts: str | list[str], ) -> np.ndarray: if isinstance(texts, str): texts = [texts] return np.random.randn(len(texts), 4).astype(np.float32) index.embed_all(fake_embed, "test", 4) env = TreeEnvironment(index) def query_embed( texts: str | list[str], ) -> np.ndarray: if isinstance(texts, str): texts = [texts] return np.ones((len(texts), 4), dtype=np.float32) * 0.5 results = env.search_similar("test", top_k=2, embed_fn=query_embed) assert len(results) > 0 def test_no_embed_fn_raises(self) -> None: """未提供 embed_fn 时应报错。""" index = _make_test_index() env = TreeEnvironment(index) with pytest.raises(ValueError, match="embed_fn"): env.search_similar("test", top_k=3) class TestGetSubtitle: """get_subtitle 方法测试。""" def test_existing_subtitle(self) -> None: """有字幕的节点应返回字幕文本。""" env = TreeEnvironment(_make_test_index()) assert env.get_subtitle("vid_L1_000_L2_000_L3_000") == "he is running" def test_no_subtitle(self) -> None: """无字幕的节点应返回空字符串。""" env = TreeEnvironment(_make_test_index()) assert env.get_subtitle("vid_L1_000_L2_000_L3_001") == "" def test_unknown_node(self) -> None: """不存在的节点应返回空字符串。""" env = TreeEnvironment(_make_test_index()) assert env.get_subtitle("nonexistent") == "" class TestResolveFramePaths: """resolve_frame_paths 方法测试。""" def test_l3_nodes(self) -> None: """L3 节点应映射到帧文件路径。""" env = TreeEnvironment(_make_test_index(), frames_dir=Path("/data/frames")) paths = env.resolve_frame_paths(["vid_L1_000_L2_000_L3_000"]) assert len(paths) == 1 assert "L1_000_L2_000_L3_000" in str(paths[0]) def test_l2_expands_to_children(self) -> None: """L2 节点应展开为其所有 L3 子节点的帧路径。""" env = TreeEnvironment(_make_test_index(), frames_dir=Path("/data/frames")) paths = env.resolve_frame_paths(["vid_L1_000_L2_000"]) assert len(paths) == 2 # 2 L3 children def test_no_frames_dir_uses_node_path(self) -> None: """未提供 frames_dir 时应使用节点自带的 frame_path。""" env = TreeEnvironment(_make_test_index()) paths = env.resolve_frame_paths(["vid_L1_000_L2_000_L3_000"]) assert len(paths) == 1 assert "L1_000_L2_000_L3_000" in str(paths[0]) def test_empty_list_returns_empty(self) -> None: """空列表应返回空结果。""" env = TreeEnvironment(_make_test_index(), frames_dir=Path("/data/frames")) paths = env.resolve_frame_paths([]) assert paths == [] class TestGetNodeText: """get_node_text 方法测试。""" def test_normal_mode_returns_full_text(self) -> None: """默认模式应返回完整文本和 None anchor_map。""" env = TreeEnvironment(_make_test_index()) text, anchor_map = env.get_node_text("vid_L1_000_L2_000_L3_000") assert "运动员在跑步" in text assert anchor_map is None def test_anchor_mode_returns_anchored_text_and_map(self) -> None: """锚模式应返回带锚文本和 anchor_map 字典。""" env = TreeEnvironment(_make_test_index()) text, anchor_map = env.get_node_text( "vid_L1_000_L2_000_L3_000", anchor=True, ) # 锚文本包含 [cN] 标记 assert "[c1]" in text # anchor_map 非空,键为锚标(如 "c1"),值为对应行文本 assert anchor_map is not None assert len(anchor_map) > 0 assert "c1" in anchor_map # 字幕行也应在 anchor_map 中(该节点有 subtitle) assert any(k.startswith("s") for k in anchor_map) def test_nonexistent_node_raises(self) -> None: """查询不存在的节点应抛出 KeyError。""" env = TreeEnvironment(_make_test_index()) with pytest.raises(KeyError): env.get_node_text("nonexistent") def test_node_without_subtitle_no_s_anchors(self) -> None: """无字幕的 L3 节点锚模式不应产生 [sN] 锚。""" env = TreeEnvironment(_make_test_index()) text, anchor_map = env.get_node_text( "vid_L1_000_L2_000_L3_001", anchor=True, ) assert anchor_map is not None assert not any(k.startswith("s") for k in anchor_map) class TestGetChildrenInfo: """get_children_info 方法测试。""" def test_l1_has_children(self) -> None: """L1 节点应返回其 L2 子节点信息列表。""" env = TreeEnvironment(_make_test_index()) children = env.get_children_info("vid_L1_000") assert len(children) == 1 child = children[0] assert child["id"] == "vid_L1_000_L2_000" assert "time_range" in child assert "summary" in child assert isinstance(child["summary"], str) def test_l2_has_children(self) -> None: """L2 节点应返回其 L3 子节点信息列表。""" env = TreeEnvironment(_make_test_index()) children = env.get_children_info("vid_L1_000_L2_000") assert len(children) == 2 ids = [c["id"] for c in children] assert "vid_L1_000_L2_000_L3_000" in ids assert "vid_L1_000_L2_000_L3_001" in ids def test_l3_has_no_children(self) -> None: """L3 叶子节点应返回空列表。""" env = TreeEnvironment(_make_test_index()) children = env.get_children_info("vid_L1_000_L2_000_L3_000") assert children == [] def test_nonexistent_node_raises(self) -> None: """查询不存在的节点应抛出 KeyError。""" env = TreeEnvironment(_make_test_index()) with pytest.raises(KeyError): env.get_children_info("nonexistent") def test_summary_truncation(self) -> None: """超过 120 字符的描述应被截断。""" index = _make_test_index() # 修改 L2 的事件描述为超长文本 l2 = index.roots[0].children[0] long_desc = "A" * 200 object.__setattr__(l2.card, "event_description", long_desc) env = TreeEnvironment(index) children = env.get_children_info("vid_L1_000") assert len(children[0]["summary"]) == 123 # 120 + "..."