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

"""修复重生成器单元测试。"""

from __future__ import annotations

import asyncio
import json
from pathlib import Path

from app.tree.index import (
    IndexMeta,
    L1Card,
    L1Node,
    L2Card,
    L2Node,
    L3Card,
    L3Node,
    TreeIndex,
)
from app.tree.repair.detector import NodeIssue
from app.tree.repair.regenerator import RepairStats, repair_tree
from core.types import LLMResponse


def _mock_response(content: str) -> LLMResponse:
    """构造模拟 LLMResponse。"""
    return LLMResponse(
        content=content,
        thinking="",
        model="mock",
        provider="mock",
        prompt_tokens=0,
        completion_tokens=0,
        latency_ms=0,
        ttft_ms=None,
        max_inter_token_ms=None,
        cache_hit=False,
        call_id="mock",
    )


class MockVLM:
    """模拟 VLM 端口，返回固定的 L3Card JSON。"""

    def __init__(self) -> None:
        self.call_count = 0

    async def chat_with_images(
        self,
        messages: list[dict],
        images: list,
        **kw: object,
    ) -> LLMResponse:
        """模拟 VLM 图文调用。"""
        self.call_count += 1
        return _mock_response(
            json.dumps(
                {
                    "frame_summary": "修复后的帧描述",
                    "visible_entities": ["修复实体"],
                    "ongoing_actions": ["修复动作"],
                    "visible_text": [],
                    "spatial_layout": "居中",
                    "visual_attributes": {"lighting": "明亮"},
                }
            )
        )


class MockLLM:
    """模拟 LLM 端口，根据 prompt 内容返回 L2Card 或 L1Card JSON。"""

    def __init__(self) -> None:
        self.call_count = 0

    async def chat(
        self,
        messages: list[dict],
        **kw: object,
    ) -> LLMResponse:
        """模拟 LLM 文本调用，按 prompt 内容区分 L2/L1 响应。"""
        self.call_count += 1
        content = messages[-1].get("content", "")
        if "段落" in content or "scene" in content.lower():
            return _mock_response(
                json.dumps(
                    {
                        "scene_summary": "修复后的场景",
                        "main_setting": "室内",
                        "key_entities": [],
                        "main_actions": [],
                        "topic_keywords": [],
                        "visible_text": [],
                        "temporal_flow": "",
                    }
                )
            )
        return _mock_response(
            json.dumps(
                {
                    "event_description": "修复后的事件",
                    "entities": [],
                    "actions": [],
                    "action_subjects": [],
                    "visible_text": [],
                    "spatial_relations": "",
                    "state_changes": None,
                }
            )
        )


class TestRepairTree:
    """repair_tree 核心测试。"""

    def _make_broken_tree(
        self,
        tmp_path: Path,
    ) -> tuple[TreeIndex, list[NodeIssue]]:
        """构建含一个空 frame_summary 的 L3 节点的测试树。"""
        frame_path = tmp_path / "frames" / "L1_000_L2_000_L3_000.jpg"
        frame_path.parent.mkdir(parents=True)
        frame_path.write_bytes(b"\xff\xd8\xff\xe0fake")

        l3 = L3Node(
            id="vid_L1_000_L2_000_L3_000",
            card=L3Card("", [], [], [], "", {}),
            timestamp=1.0,
            frame_path="frames/L1_000_L2_000_L3_000.jpg",
        )
        l2 = L2Node(
            id="vid_L1_000_L2_000",
            card=L2Card("原始事件", [], [], [], [], "", None),
            time_range=(0.0, 10.0),
            children=[l3],
        )
        l1 = L1Node(
            id="vid_L1_000",
            card=L1Card("原始场景", "", [], [], [], [], ""),
            time_range=(0.0, 10.0),
            children=[l2],
        )
        index = TreeIndex(metadata=IndexMeta("/t.mp4", "video"), roots=[l1])
        issues = [
            NodeIssue(
                "vid_L1_000_L2_000_L3_000",
                3,
                "empty_field",
                "frame_summary 为空",
            )
        ]
        return index, issues

    def test_repairs_l3_and_cascades(self, tmp_path: Path) -> None:
        """修复 L3 后应级联重生成 L2 和 L1。"""
        index, issues = self._make_broken_tree(tmp_path)
        stats = asyncio.run(repair_tree(index, issues, MockVLM(), MockLLM(), tmp_path))
        assert stats.l3_repaired == 1
        assert stats.l2_regenerated == 1
        assert stats.l1_regenerated == 1
        assert index.roots[0].children[0].children[0].card.frame_summary == "修复后的帧描述"
        assert index.roots[0].children[0].card.event_description == "修复后的事件"
        assert index.roots[0].card.scene_summary == "修复后的场景"

    def test_no_issues_no_changes(self) -> None:
        """无问题时不进行任何修复。"""
        l3 = L3Node(
            id="l1_0_l2_0_l3_0",
            card=L3Card("正常", [], [], [], "", {}),
            timestamp=1.0,
        )
        l2 = L2Node(
            id="l1_0_l2_0",
            card=L2Card("正常事件", [], [], [], [], "", None),
            time_range=(0.0, 10.0),
            children=[l3],
        )
        l1 = L1Node(
            id="l1_0",
            card=L1Card("正常场景", "", [], [], [], [], ""),
            time_range=(0.0, 10.0),
            children=[l2],
        )
        index = TreeIndex(metadata=IndexMeta("/t.mp4", "video"), roots=[l1])
        stats = asyncio.run(repair_tree(index, [], MockVLM(), MockLLM(), Path("/tmp")))
        assert stats.l3_repaired == 0
        assert stats.l2_regenerated == 0
        assert stats.l1_regenerated == 0

    def test_stats_dataclass(self) -> None:
        """RepairStats 数据类字段验证。"""
        stats = RepairStats(l3_repaired=2, l2_regenerated=1, l1_regenerated=1)
        assert stats.l3_repaired == 2
        assert stats.l2_regenerated == 1
        assert stats.l1_regenerated == 1

    def test_skips_non_empty_field_issues(self, tmp_path: Path) -> None:
        """非 empty_field 类型的 issue 不触发 L3 修复。"""
        l3 = L3Node(
            id="l1_0_l2_0_l3_0",
            card=L3Card("正常描述", [], [], [], "", {}),
            timestamp=1.0,
        )
        l2 = L2Node(
            id="l1_0_l2_0",
            card=L2Card("原始事件", [], [], [], [], "", None),
            time_range=(0.0, 10.0),
            children=[l3],
        )
        l1 = L1Node(
            id="l1_0",
            card=L1Card("原始场景", "", [], [], [], [], ""),
            time_range=(0.0, 10.0),
            children=[l2],
        )
        index = TreeIndex(metadata=IndexMeta("/t.mp4", "video"), roots=[l1])
        issues = [NodeIssue("l1_0_l2_0_l3_0", 3, "missing_frame", "帧文件不存在")]
        stats = asyncio.run(repair_tree(index, issues, MockVLM(), MockLLM(), tmp_path))
        assert stats.l3_repaired == 0
        assert stats.l2_regenerated == 0

    def test_multiple_l3_under_same_l2(self, tmp_path: Path) -> None:
        """同一 L2 下多个 L3 修复后，L2 只重生成一次。"""
        frame_dir = tmp_path / "frames"
        frame_dir.mkdir(parents=True)
        for i in range(2):
            (frame_dir / f"f{i}.jpg").write_bytes(b"\xff\xd8\xff\xe0fake")

        l3_a = L3Node(
            id="l1_0_l2_0_l3_0",
            card=L3Card("", [], [], [], "", {}),
            timestamp=1.0,
            frame_path="frames/f0.jpg",
        )
        l3_b = L3Node(
            id="l1_0_l2_0_l3_1",
            card=L3Card("", [], [], [], "", {}),
            timestamp=2.0,
            frame_path="frames/f1.jpg",
        )
        l2 = L2Node(
            id="l1_0_l2_0",
            card=L2Card("原始事件", [], [], [], [], "", None),
            time_range=(0.0, 10.0),
            children=[l3_a, l3_b],
        )
        l1 = L1Node(
            id="l1_0",
            card=L1Card("原始场景", "", [], [], [], [], ""),
            time_range=(0.0, 10.0),
            children=[l2],
        )
        index = TreeIndex(metadata=IndexMeta("/t.mp4", "video"), roots=[l1])
        issues = [
            NodeIssue("l1_0_l2_0_l3_0", 3, "empty_field", "frame_summary 为空"),
            NodeIssue("l1_0_l2_0_l3_1", 3, "empty_field", "frame_summary 为空"),
        ]
        vlm = MockVLM()
        llm = MockLLM()
        stats = asyncio.run(repair_tree(index, issues, vlm, llm, tmp_path))
        assert stats.l3_repaired == 2
        assert stats.l2_regenerated == 1
        assert stats.l1_regenerated == 1
        assert vlm.call_count == 2
        # LLM 应被调用 2 次：一次 L2 + 一次 L1
        assert llm.call_count == 2

    def test_missing_frame_file_skips_l3(self, tmp_path: Path) -> None:
        """帧文件不存在时跳过该 L3 节点的修复。"""
        l3 = L3Node(
            id="l1_0_l2_0_l3_0",
            card=L3Card("", [], [], [], "", {}),
            timestamp=1.0,
            frame_path="frames/nonexistent.jpg",
        )
        l2 = L2Node(
            id="l1_0_l2_0",
            card=L2Card("原始事件", [], [], [], [], "", None),
            time_range=(0.0, 10.0),
            children=[l3],
        )
        l1 = L1Node(
            id="l1_0",
            card=L1Card("原始场景", "", [], [], [], [], ""),
            time_range=(0.0, 10.0),
            children=[l2],
        )
        index = TreeIndex(metadata=IndexMeta("/t.mp4", "video"), roots=[l1])
        issues = [NodeIssue("l1_0_l2_0_l3_0", 3, "empty_field", "frame_summary 为空")]
        stats = asyncio.run(repair_tree(index, issues, MockVLM(), MockLLM(), tmp_path))
        # 帧文件不存在 → 跳过 L3 修复 → 无级联
        assert stats.l3_repaired == 0
        assert stats.l2_regenerated == 0
        assert stats.l1_regenerated == 0