feat(tree): add get_node_text + get_children_info to TreeEnvironment

- get_node_text(node_id, anchor=False): returns raw text + optional
  anchor_map dict by parsing [cN]/[sN] prefixes from anchored text
- get_children_info(node_id): returns structured child list with
  id/time_range/summary (description truncated to 120 chars)
- Both methods reuse existing internal helpers (_node_full_text,
  _node_anchored_text, _get_children, _node_description,
  _format_time_range)
- 9 new test cases across TestGetNodeText and TestGetChildrenInfo

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-07-07 05:45:48 -04:00
parent ffb37a1c1e
commit 44ee62867d
2 changed files with 843 additions and 0 deletions
+324
View File
@@ -0,0 +1,324 @@
"""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 + "..."