"""TreeIndex 数据结构单元测试。""" from __future__ import annotations import json import numpy as np import pytest from app.tree.index import ( IndexMeta, L1Card, L1Node, L2Card, L2Node, L3Card, L3Node, TreeIndex, ) def _make_l3(idx: int = 0) -> L3Node: return L3Node( id=f"l1_0_l2_0_l3_{idx}", card=L3Card( frame_summary=f"帧{idx}描述", visible_entities=["实体A"], ongoing_actions=["动作A"], visible_text=["文字A"], spatial_layout="居中构图", visual_attributes={"lighting": "明亮"}, ), timestamp=idx * 2.0, frame_path=f"frames/l1_0_l2_0_l3_{idx}.jpg", ) def _make_l2(n_l3: int = 2) -> L2Node: return L2Node( id="l1_0_l2_0", card=L2Card( event_description="事件描述", entities=["实体B"], actions=["动作B"], action_subjects=["主体B"], visible_text=["文字B"], spatial_relations="左右排列", state_changes=None, ), time_range=(0.0, 60.0), children=[_make_l3(i) for i in range(n_l3)], ) def _make_l1(n_l2: int = 1, n_l3: int = 2) -> L1Node: return L1Node( id="l1_0", card=L1Card( scene_summary="场景摘要", main_setting="室内", key_entities=["实体C"], main_actions=["动作C"], topic_keywords=["关键词"], visible_text=["文字C"], temporal_flow="从左到右", ), time_range=(0.0, 600.0), children=[_make_l2(n_l3) for _ in range(n_l2)], ) def _make_index(n_l1: int = 1) -> TreeIndex: meta = IndexMeta(source_path="/test/video.mp4", modality="video") return TreeIndex(metadata=meta, roots=[_make_l1() for _ in range(n_l1)]) class TestCards: def test_l3_card_frozen(self): card = L3Card( frame_summary="desc", visible_entities=[], ongoing_actions=[], visible_text=[], spatial_layout="", visual_attributes={}, ) with pytest.raises(AttributeError): card.frame_summary = "changed" def test_l2_card_fields(self): card = L2Card( event_description="evt", entities=[], actions=[], action_subjects=[], visible_text=[], spatial_relations="", state_changes=None, ) assert card.event_description == "evt" assert card.state_changes is None def test_l1_card_fields(self): card = L1Card( scene_summary="scene", main_setting="outdoor", key_entities=["e"], main_actions=["a"], topic_keywords=["k"], visible_text=["t"], temporal_flow="flow", ) assert card.scene_summary == "scene" class TestNodes: def test_l3_description_property(self): node = _make_l3() assert node.description == node.card.frame_summary def test_l2_description_property(self): node = _make_l2() assert node.description == node.card.event_description def test_l1_summary_property(self): node = _make_l1() assert node.summary == node.card.scene_summary def test_l3_default_embedding_none(self): node = _make_l3() assert node.embedding is None def test_l3_subtitle_default_none(self): node = _make_l3() assert node.subtitle is None class TestTreeIndex: def test_is_embedded_false_by_default(self): index = _make_index() assert not index.is_embedded def test_embed_all(self): index = _make_index() def fake_embed(texts): if isinstance(texts, str): texts = [texts] return np.random.randn(len(texts), 4).astype(np.float32) index.embed_all(fake_embed, "test-model", 4) assert index.is_embedded assert index.metadata.embed_model == "test-model" assert index.metadata.embed_dim == 4 def test_l1_embeddings_shape(self): index = _make_index(n_l1=2) def fake_embed(texts): if isinstance(texts, str): texts = [texts] return np.random.randn(len(texts), 4).astype(np.float32) index.embed_all(fake_embed, "test-model", 4) m = index.l1_embeddings() assert m.shape == (2, 4) def test_get_node(self): index = _make_index() node = index.get_node(0, 0, 1) assert node.id == "l1_0_l2_0_l3_1" def test_get_node_out_of_bounds(self): index = _make_index() with pytest.raises(IndexError): index.get_node(99, 0, 0) class TestSerialization: def test_json_roundtrip(self, tmp_path): index = _make_index() 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].id == "l1_0" assert loaded.roots[0].card.scene_summary == "场景摘要" assert loaded.roots[0].children[0].children[0].card.frame_summary == "帧0描述" def test_json_roundtrip_with_embedding(self, tmp_path): index = _make_index() def fake_embed(texts): if isinstance(texts, str): texts = [texts] return np.random.randn(len(texts), 4).astype(np.float32) index.embed_all(fake_embed, "test-model", 4) path = tmp_path / "tree_emb.json" index.save_json(str(path), include_embedding=True) loaded = TreeIndex.load_json(str(path)) assert loaded.is_embedded np.testing.assert_array_almost_equal( loaded.roots[0].embedding, index.roots[0].embedding, decimal=5 ) def test_l1_json_roundtrip(self, tmp_path): from app.tree.index import load_l1_json, save_l1_json l1 = _make_l1() path = tmp_path / "l1_0.json" save_l1_json(str(path), l1) loaded = load_l1_json(str(path)) assert loaded.id == "l1_0" assert len(loaded.children) == 1 assert len(loaded.children[0].children) == 2 def test_id_uniqueness_validation(self, tmp_path): """重复 ID 在反序列化时应报错。""" index = _make_index() d = index.to_dict() d["roots"].append(d["roots"][0]) path = tmp_path / "dup.json" with open(path, "w") as f: json.dump(d, f) with pytest.raises(ValueError, match="重复"): TreeIndex.load_json(str(path))