Video-Tree-TRM5/tests/unit/test_tree_index.py

"""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))