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

"""observe_frame 单元测试。

FakeVLMProvider / FakeOCRProvider 实现 Protocol 最小集，
覆盖：两轮正常、verify=False 仅提取、OCR 注入、OCR 失败降级、
OCR 为 None、VLM 提取失败、VLM 验证失败降级、帧文件不存在、stats 键完整性。
"""

from __future__ import annotations

import asyncio
from pathlib import Path
from typing import Any

import pytest

from core.types import LLMResponse

# ---------------------------------------------------------------------------
# Fake 实现
# ---------------------------------------------------------------------------

_DUMMY_RESPONSE_KWARGS = {
    "thinking": "",
    "model": "fake-vlm",
    "provider": "fake",
    "prompt_tokens": 10,
    "completion_tokens": 20,
    "latency_ms": 100,
    "ttft_ms": None,
    "max_inter_token_ms": None,
    "cache_hit": False,
    "call_id": "fake-call-id",
}


class FakeVLMProvider:
    """可编程的 VLM 假实现。

    通过 responses 列表按序返回预设内容；raises 列表对应位置不为 None 时抛异常。
    """

    def __init__(
        self,
        responses: list[str] | None = None,
        raises: list[Exception | None] | None = None,
    ) -> None:
        self._responses = responses or []
        self._raises = raises or []
        self._call_idx = 0
        self.calls: list[dict[str, Any]] = []

    async def chat_with_images(
        self,
        messages: list[dict[str, Any]],
        images: list[str | Path],
        *,
        session_id: str | None = None,
        parent_call_id: str | None = None,
    ) -> LLMResponse:
        """记录调用并按序返回预设响应或抛出异常。"""
        idx = self._call_idx
        self._call_idx += 1
        self.calls.append(
            {
                "messages": messages,
                "images": images,
                "session_id": session_id,
                "parent_call_id": parent_call_id,
            }
        )
        if idx < len(self._raises) and self._raises[idx] is not None:
            raise self._raises[idx]  # type: ignore[misc]
        content = self._responses[idx] if idx < len(self._responses) else ""
        return LLMResponse(content=content, **_DUMMY_RESPONSE_KWARGS)


class FakeOCRProvider:
    """可编程的 OCR 假实现。"""

    def __init__(
        self,
        text: str = "",
        raise_on_call: Exception | None = None,
    ) -> None:
        self._text = text
        self._raise_on_call = raise_on_call

    async def transcribe_frames(self, frame_paths: list[Path]) -> str:
        """返回预设文本或抛出异常。"""
        if self._raise_on_call is not None:
            raise self._raise_on_call
        return self._text


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture()
def frame_files(tmp_path: Path) -> list[Path]:
    """创建两个最小 JPEG 占位帧文件。"""
    frames: list[Path] = []
    for i in range(2):
        p = tmp_path / f"frame_{i}.jpg"
        p.write_bytes(b"\xff\xd8\xff\xe0" + b"\x00" * 20)
        frames.append(p)
    return frames


@pytest.fixture()
def prompts_dir(tmp_path: Path) -> Path:
    """创建 observe_frame_extract.md / observe_frame_verify.md 占位文件。"""
    d = tmp_path / "prompts"
    d.mkdir()
    (d / "observe_frame_extract.md").write_text("extract prompt", encoding="utf-8")
    (d / "observe_frame_verify.md").write_text("verify prompt", encoding="utf-8")
    return d


# ---------------------------------------------------------------------------
# 测试用例
# ---------------------------------------------------------------------------


class TestObserveFrameNormal:
    """两轮正常执行（verify=True）。"""

    def test_two_round_normal(
        self, frame_files: list[Path], prompts_dir: Path
    ) -> None:
        from app.search.vision import observe_frame

        vlm = FakeVLMProvider(responses=["raw evidence", "verified ok"])
        collected: list[dict[str, int]] = []

        result = asyncio.run(
            observe_frame(
                vlm=vlm,
                frame_paths=frame_files,
                question="what happened?",
                prompts_dir=prompts_dir,
                ocr=None,
                verify=True,
                stats_sink=collected.append,
            )
        )

        assert result == "[视觉观察] raw evidence\n[验证] verified ok"
        assert len(vlm.calls) == 2
        # 提取轮使用 extract prompt
        assert vlm.calls[0]["messages"][0]["content"] == "extract prompt"
        # 验证轮使用 verify prompt
        assert vlm.calls[1]["messages"][0]["content"] == "verify prompt"
        assert len(collected) == 1


class TestObserveFrameExtractOnly:
    """verify=False 仅执行提取轮。"""

    def test_extract_only(
        self, frame_files: list[Path], prompts_dir: Path
    ) -> None:
        from app.search.vision import observe_frame

        vlm = FakeVLMProvider(responses=["only extract"])

        result = asyncio.run(
            observe_frame(
                vlm=vlm,
                frame_paths=frame_files,
                question="q?",
                prompts_dir=prompts_dir,
                ocr=None,
                verify=False,
            )
        )

        assert result == "[视觉观察] only extract"
        assert len(vlm.calls) == 1


class TestObserveFrameOCRInjection:
    """OCR 注入：文本非空时并置于问题前。"""

    def test_ocr_injected(
        self, frame_files: list[Path], prompts_dir: Path
    ) -> None:
        from app.search.vision import observe_frame

        vlm = FakeVLMProvider(responses=["evidence with ocr"])
        ocr = FakeOCRProvider(text="帧1: 你好世界")
        collected: list[dict[str, int]] = []

        result = asyncio.run(
            observe_frame(
                vlm=vlm,
                frame_paths=frame_files,
                question="q?",
                prompts_dir=prompts_dir,
                ocr=ocr,
                verify=False,
                stats_sink=collected.append,
            )
        )

        assert "[视觉观察]" in result
        # 验证 OCR 文本被注入到 user message
        user_msg = vlm.calls[0]["messages"][1]["content"]
        assert "帧1: 你好世界" in user_msg
        # stats 中 ocr_injected=1
        assert collected[0]["ocr_injected"] == 1
        assert collected[0]["ocr_chars"] == len("帧1: 你好世界")


class TestObserveFrameOCRFailDegrades:
    """OCR 转录抛出异常时降级：不注入 OCR、ocr_failed=1。"""

    def test_ocr_failure_degrades(
        self, frame_files: list[Path], prompts_dir: Path
    ) -> None:
        from app.search.vision import observe_frame

        vlm = FakeVLMProvider(responses=["evidence no ocr"])
        ocr = FakeOCRProvider(raise_on_call=RuntimeError("OCR service down"))
        collected: list[dict[str, int]] = []

        result = asyncio.run(
            observe_frame(
                vlm=vlm,
                frame_paths=frame_files,
                question="q?",
                prompts_dir=prompts_dir,
                ocr=ocr,
                verify=False,
                stats_sink=collected.append,
            )
        )

        assert result == "[视觉观察] evidence no ocr"
        assert collected[0]["ocr_failed"] == 1
        assert collected[0]["ocr_injected"] == 0


class TestObserveFrameOCRNone:
    """ocr=None 时不执行转录。"""

    def test_ocr_none(
        self, frame_files: list[Path], prompts_dir: Path
    ) -> None:
        from app.search.vision import observe_frame

        vlm = FakeVLMProvider(responses=["no ocr"])
        collected: list[dict[str, int]] = []

        result = asyncio.run(
            observe_frame(
                vlm=vlm,
                frame_paths=frame_files,
                question="q?",
                prompts_dir=prompts_dir,
                ocr=None,
                verify=False,
                stats_sink=collected.append,
            )
        )

        assert result == "[视觉观察] no ocr"
        assert collected[0]["ocr_injected"] == 0
        assert collected[0]["ocr_chars"] == 0
        assert collected[0]["ocr_failed"] == 0


class TestObserveFrameVLMExtractFailure:
    """VLM 提取轮失败 → 返回 [VL错误]。"""

    def test_vlm_extract_failure(
        self, frame_files: list[Path], prompts_dir: Path
    ) -> None:
        from app.search.vision import observe_frame

        vlm = FakeVLMProvider(raises=[RuntimeError("VLM timeout")])
        collected: list[dict[str, int]] = []

        result = asyncio.run(
            observe_frame(
                vlm=vlm,
                frame_paths=frame_files,
                question="q?",
                prompts_dir=prompts_dir,
                ocr=None,
                verify=True,
                stats_sink=collected.append,
            )
        )

        assert result.startswith("[VL错误]")
        assert "VLM timeout" in result
        assert len(collected) == 1


class TestObserveFrameVLMVerifyFailureDegrades:
    """VLM 验证轮失败 → 降级：保留提取结果 + [验证] 跳过。"""

    def test_vlm_verify_failure_degrades(
        self, frame_files: list[Path], prompts_dir: Path
    ) -> None:
        from app.search.vision import observe_frame

        vlm = FakeVLMProvider(
            responses=["good evidence", ""],
            raises=[None, RuntimeError("verify timeout")],
        )
        collected: list[dict[str, int]] = []

        result = asyncio.run(
            observe_frame(
                vlm=vlm,
                frame_paths=frame_files,
                question="q?",
                prompts_dir=prompts_dir,
                ocr=None,
                verify=True,
                stats_sink=collected.append,
            )
        )

        assert "[视觉观察] good evidence" in result
        assert "[验证] 跳过（调用失败）" in result
        assert len(collected) == 1


class TestObserveFrameFileMissing:
    """帧文件不存在 → 返回 [VL错误] 帧文件不存在。"""

    def test_frame_file_not_found(self, prompts_dir: Path) -> None:
        from app.search.vision import observe_frame

        vlm = FakeVLMProvider()
        missing = [Path("/nonexistent/frame_0.jpg")]
        collected: list[dict[str, int]] = []

        result = asyncio.run(
            observe_frame(
                vlm=vlm,
                frame_paths=missing,
                question="q?",
                prompts_dir=prompts_dir,
                ocr=None,
                verify=True,
                stats_sink=collected.append,
            )
        )

        assert "[VL错误] 帧文件不存在" in result
        assert len(collected) == 1


class TestObserveFrameStatsKeys:
    """stats 包含全部五个预期键。"""

    def test_stats_keys_complete(
        self, frame_files: list[Path], prompts_dir: Path
    ) -> None:
        from app.search.vision import observe_frame

        vlm = FakeVLMProvider(responses=["evidence"])
        collected: list[dict[str, int]] = []

        asyncio.run(
            observe_frame(
                vlm=vlm,
                frame_paths=frame_files,
                question="q?",
                prompts_dir=prompts_dir,
                ocr=None,
                verify=False,
                stats_sink=collected.append,
            )
        )

        expected_keys = {"ocr_injected", "ocr_chars", "ocr_failed", "discrepancy", "abstain"}
        assert set(collected[0].keys()) == expected_keys


class TestObserveFrameDiscrepancyAndAbstain:
    """VLM 返回含 '分歧' 或 '[证据不存在]' 时对应 stats 标记。"""

    def test_discrepancy_flag(
        self, frame_files: list[Path], prompts_dir: Path
    ) -> None:
        from app.search.vision import observe_frame

        vlm = FakeVLMProvider(responses=["发现分歧：OCR 与画面不一致"])
        collected: list[dict[str, int]] = []

        asyncio.run(
            observe_frame(
                vlm=vlm,
                frame_paths=frame_files,
                question="q?",
                prompts_dir=prompts_dir,
                ocr=None,
                verify=False,
                stats_sink=collected.append,
            )
        )

        assert collected[0]["discrepancy"] == 1

    def test_abstain_flag(
        self, frame_files: list[Path], prompts_dir: Path
    ) -> None:
        from app.search.vision import observe_frame

        vlm = FakeVLMProvider(responses=["[证据不存在] 无法判断"])
        collected: list[dict[str, int]] = []

        asyncio.run(
            observe_frame(
                vlm=vlm,
                frame_paths=frame_files,
                question="q?",
                prompts_dir=prompts_dir,
                ocr=None,
                verify=False,
                stats_sink=collected.append,
            )
        )

        assert collected[0]["abstain"] == 1


class TestObserveFrameTelemetryPassthrough:
    """session_id 和 parent_call_id 透传到 VLM 调用。"""

    def test_telemetry_passthrough(
        self, frame_files: list[Path], prompts_dir: Path
    ) -> None:
        from app.search.vision import observe_frame

        vlm = FakeVLMProvider(responses=["evidence"])

        asyncio.run(
            observe_frame(
                vlm=vlm,
                frame_paths=frame_files,
                question="q?",
                prompts_dir=prompts_dir,
                ocr=None,
                verify=False,
                session_id="sess-123",
                parent_call_id="parent-456",
            )
        )

        assert vlm.calls[0]["session_id"] == "sess-123"
        assert vlm.calls[0]["parent_call_id"] == "parent-456"