iomgaa
44ee62867d
- 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>
325 lines
11 KiB
Python
325 lines
11 KiB
Python
"""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 + "..."
|