Video-Tree-TRM5/reference/tests/unit/test_embedding_persistence.py

"""
Embedding 持久化单元测试
=======================
测试 TreeIndex 的 embedding 序列化/反序列化功能。
"""

import json
import os
import tempfile
from pathlib import Path

import numpy as np
import pytest

from video_tree_trm.tree_index import (
    L1Node,
    L2Node,
    L3Node,
    IndexMeta,
    TreeIndex,
    _embed_to_str,
    _embed_from_str,
)


class TestEmbeddingSerialization:
    """测试 embedding 序列化辅助函数。"""

    def test_embed_to_str_and_back(self):
        """测试 base64 序列化/反序列化往返正确。"""
        arr = np.random.randn(768).astype(np.float32)
        s = _embed_to_str(arr)
        recovered = _embed_from_str(s)

        assert recovered is not None
        assert recovered.dtype == np.float32
        assert recovered.shape == (768,)
        np.testing.assert_array_almost_equal(arr, recovered, decimal=6)

    def test_embed_to_str_handles_none(self):
        """测试 None 输入返回 None。"""
        assert _embed_to_str(None) is None
        assert _embed_from_str(None) is None
        assert _embed_from_str("") is None


class TestL3NodeEmbedding:
    """测试 L3Node embedding 序列化。"""

    def test_l3_to_dict_without_embedding(self):
        """测试不带 embedding 序列化。"""
        node = L3Node(
            id="l3_0",
            description="测试描述",
            timestamp=1.0,
            frame_path="frame.jpg",
            raw_content="原始内容",
        )
        d = {
            "id": node.id,
            "description": node.description,
            "timestamp": node.timestamp,
            "frame_path": node.frame_path,
            "raw_content": node.raw_content,
        }
        # 无 embedding 字段
        assert "embedding" not in d

    def test_l3_embedding_roundtrip(self):
        """测试 L3 embedding 序列化往返。"""
        embed = np.random.randn(768).astype(np.float32)
        s = _embed_to_str(embed)
        recovered = _embed_from_str(s)
        np.testing.assert_array_almost_equal(embed, recovered, decimal=6)


class TestTreeIndexEmbedding:
    """测试 TreeIndex 完整序列化/反序列化。"""

    def test_save_load_without_embedding(self):
        """测试不带 embedding 保存/加载（向后兼容）。"""
        with tempfile.TemporaryDirectory() as tmpdir:
            path = os.path.join(tmpdir, "test.json")

            # 创建简单树
            l3 = L3Node(id="l3_0", description="L3 描述")
            l2 = L2Node(id="l2_0", description="L2 描述", children=[l3])
            l1 = L1Node(id="l1_0", summary="L1 摘要", children=[l2])
            meta = IndexMeta(source_path="test.mp4", modality="video")
            tree = TreeIndex(metadata=meta, roots=[l1])

            # 保存（不含 embedding）
            tree.save_json(path, include_embedding=False)

            # 加载
            loaded = TreeIndex.load_json(path)

            assert len(loaded.roots) == 1
            assert loaded.roots[0].summary == "L1 摘要"
            assert not loaded.is_embedded  # 无 embedding

    def test_save_load_with_embedding(self):
        """测试带 embedding 保存/加载。"""
        with tempfile.TemporaryDirectory() as tmpdir:
            path = os.path.join(tmpdir, "test_embed.json")

            # 创建带 embedding 的树
            embed_l1 = np.random.randn(768).astype(np.float32)
            embed_l2 = np.random.randn(768).astype(np.float32)
            embed_l3 = np.random.randn(768).astype(np.float32)

            l3 = L3Node(id="l3_0", description="L3 描述", embedding=embed_l3)
            l2 = L2Node(id="l2_0", description="L2 描述", embedding=embed_l2, children=[l3])
            l1 = L1Node(id="l1_0", summary="L1 摘要", embedding=embed_l1, children=[l2])
            meta = IndexMeta(
                source_path="test.mp4",
                modality="video",
                embed_model="test-model",
                embed_dim=768,
            )
            tree = TreeIndex(metadata=meta, roots=[l1])

            # 保存（含 embedding）
            tree.save_json(path, include_embedding=True)

            # 验证 JSON 中有 embedding 字段
            with open(path, encoding="utf-8") as f:
                data = json.load(f)
            assert "embedding" in data["roots"][0]
            assert data["metadata"]["embed_model"] == "test-model"

            # 加载
            loaded = TreeIndex.load_json(path)

            assert loaded.is_embedded
            np.testing.assert_array_almost_equal(
                loaded.roots[0].embedding, embed_l1, decimal=6
            )
            np.testing.assert_array_almost_equal(
                loaded.roots[0].children[0].embedding, embed_l2, decimal=6
            )
            np.testing.assert_array_almost_equal(
                loaded.roots[0].children[0].children[0].embedding, embed_l3, decimal=6
            )

    def test_load_old_format_compatible(self):
        """测试加载旧格式（无 embedding 字段）JSON 兼容。"""
        with tempfile.TemporaryDirectory() as tmpdir:
            path = os.path.join(tmpdir, "old_format.json")

            # 手动创建旧格式 JSON（无 embedding 字段）
            old_data = {
                "metadata": {
                    "source_path": "test.mp4",
                    "modality": "video",
                    "created_at": "2024-01-01T00:00:00",
                },
                "roots": [
                    {
                        "id": "l1_0",
                        "summary": "L1 摘要",
                        "time_range": None,
                        "children": [
                            {
                                "id": "l2_0",
                                "description": "L2 描述",
                                "time_range": None,
                                "children": [
                                    {
                                        "id": "l3_0",
                                        "description": "L3 描述",
                                        "timestamp": 1.0,
                                        "frame_path": None,
                                        "raw_content": None,
                                    }
                                ],
                            }
                        ],
                    }
                ],
            }
            with open(path, "w", encoding="utf-8") as f:
                json.dump(old_data, f)

            # 加载
            loaded = TreeIndex.load_json(path)

            assert len(loaded.roots) == 1
            assert not loaded.is_embedded  # 无 embedding，向后兼容


if __name__ == "__main__":
    pytest.main([__file__, "-v"])