Video-Tree-TRM5/tests/integration/test_pipeline_v2.py

"""v2 出题管线集成测试 — 覆盖 slot 分配、重出循环、重量抽检与完整编排。

测试策略：使用受控 mock VLM/LLM 返回，验证管线逻辑正确性。
"""

from __future__ import annotations

import asyncio
import json
import random
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from pathlib import Path

import numpy as np
import pytest

from app.question_gen.pipeline_v2 import (
    PipelineConfig,
    PipelineResult,
    SlotAssignment,
    _assign_slots,
    _process_one_slot,
    run_pipeline_v2,
)
from app.tree.index import (
    IndexMeta,
    L1Card,
    L1Node,
    L2Card,
    L2Node,
    L3Card,
    L3Node,
    TreeIndex,
)
from core.types import LLMResponse

# ---------------------------------------------------------------------------
# 测试用 fixtures
# ---------------------------------------------------------------------------


def _make_llm_response(content: str) -> LLMResponse:
    """构造标准 LLMResponse。"""
    return LLMResponse(
        content=content,
        thinking="",
        model="test-model",
        provider="test",
        prompt_tokens=10,
        completion_tokens=20,
        latency_ms=100,
        ttft_ms=50.0,
        max_inter_token_ms=10.0,
        cache_hit=False,
        call_id="call-001",
    )


def _make_gate_pass_response() -> str:
    """门控全部通过的 JSON 响应。"""
    return json.dumps({"verdict": "pass", "reason": "looks good"})


def _make_gate_fail_response(reason: str = "quality issue") -> str:
    """门控失败的 JSON 响应。"""
    return json.dumps({"verdict": "fail", "reason": reason})


def _make_candidate_json(
    question: str = "What happened next?",
    answer: str = "A",
) -> str:
    """构造 VLM 返回的候选题 JSON。"""
    return json.dumps(
        {
            "question": question,
            "options": [
                "A. The cat jumped",
                "B. The dog ran",
                "C. Nothing happened",
                "D. It rained",
            ],
            "answer": answer,
            "difficulty": "medium",
        }
    )


def _make_tree() -> TreeIndex:
    """构造最小合法三层树。"""
    l3_nodes = [
        L3Node(
            id=f"vid_L1_000_L2_000_L3_{i:03d}",
            card=L3Card(
                frame_summary=f"Frame {i} shows activity",
                visible_entities=["person", "object"],
                ongoing_actions=["walking"],
                visible_text=[],
                spatial_layout="center",
                visual_attributes={},
                subtitle=f"Subtitle sentence {i} with some unique content here",
            ),
            timestamp=float(i * 2),
            frame_path=f"frames/L1_000_L2_000_L3_{i:03d}.jpg",
        )
        for i in range(5)
    ]
    l2 = L2Node(
        id="vid_L1_000_L2_000",
        card=L2Card(
            event_description="A person walks through the park",
            entities=["person", "park"],
            actions=["walking"],
            action_subjects=["person"],
            visible_text=[],
            spatial_relations="person in center of park",
            state_changes=None,
            subtitle="Person walking in park doing activities",
        ),
        children=l3_nodes,
    )
    l2_b = L2Node(
        id="vid_L1_000_L2_001",
        card=L2Card(
            event_description="A dog runs across the field",
            entities=["dog", "field"],
            actions=["running"],
            action_subjects=["dog"],
            visible_text=[],
            spatial_relations="dog in the field",
            state_changes=None,
            subtitle="Dog running across the field",
        ),
        children=[
            L3Node(
                id=f"vid_L1_000_L2_001_L3_{i:03d}",
                card=L3Card(
                    frame_summary=f"Dog frame {i}",
                    visible_entities=["dog"],
                    ongoing_actions=["running"],
                    visible_text=[],
                    spatial_layout="wide",
                    visual_attributes={},
                    subtitle=f"Dog subtitle {i}",
                ),
                timestamp=float(10 + i * 2),
                frame_path=f"frames/L1_000_L2_001_L3_{i:03d}.jpg",
            )
            for i in range(4)
        ],
    )
    l1 = L1Node(
        id="vid_L1_000",
        card=L1Card(
            scene_summary="Outdoor activities in a park",
            main_setting="outdoor park",
            key_entities=["person", "dog"],
            main_actions=["walking", "running"],
            topic_keywords=["outdoor", "activity"],
            visible_text=[],
            temporal_flow="sequential activities",
        ),
        children=[l2, l2_b],
    )
    meta = IndexMeta(source_path="test_video.mp4", modality="video")
    return TreeIndex(metadata=meta, roots=[l1])


class MockVLM:
    """受控 VLM mock — 自动区分生成请求和门控请求。

    检测 prompt 中是否包含 'verdict' 关键词判断请求类型：
    - 门控请求 → 返回 gate_response（默认 pass）
    - 生成请求 → 按序返回 candidate JSON
    """

    def __init__(
        self,
        responses: list[str] | None = None,
        gate_response: str | None = None,
    ) -> None:
        self._responses = responses or [_make_candidate_json()]
        self._gate_response = gate_response or _make_gate_pass_response()
        self._gen_count = 0

    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:
        prompt_text = str(messages)
        system_text = messages[0].get("content", "") if messages else ""
        if "verdict" in prompt_text.lower():
            return _make_llm_response(self._gate_response)
        if "distractor" in system_text.lower() and "grader" not in system_text.lower():
            # 候选池请求：返回 4 个 grounded 干扰项
            return _make_llm_response('{"distractors": ["蒸", "煮", "炸", "烤"]}')
        if "grader" in system_text.lower():
            # 打分请求：正解高分、3 个落区间、1 个负空间
            return _make_llm_response('{"scores": [0.90, 0.80, 0.70, 0.60, 0.20]}')
        idx = min(self._gen_count, len(self._responses) - 1)
        self._gen_count += 1
        return _make_llm_response(self._responses[idx])


class MockLLM:
    """受控 LLM mock — 支持配置门控 pass/fail 序列。"""

    def __init__(self, responses: list[str] | None = None) -> None:
        self._responses = responses or [_make_gate_pass_response()]
        self._call_count = 0

    async def chat(
        self,
        messages: list[dict[str, Any]],
        *,
        session_id: str | None = None,
        parent_call_id: str | None = None,
    ) -> LLMResponse:
        idx = min(self._call_count, len(self._responses) - 1)
        self._call_count += 1
        return _make_llm_response(self._responses[idx])


def _mock_embed_fn(text: str) -> np.ndarray:
    """确定性 embedding：基于文本 hash 生成向量。"""
    rng = np.random.default_rng(hash(text) % (2**32))
    vec = rng.standard_normal(64).astype(np.float32)
    return vec / np.linalg.norm(vec)


@pytest.fixture
def tree() -> TreeIndex:
    return _make_tree()


@pytest.fixture
def default_config(tmp_path: Path) -> PipelineConfig:
    return PipelineConfig(
        per_type=2,
        retry_limit=3,
        heavy_sample_rate=0.15,
        dedup_threshold=0.85,
        concurrency=2,
        seed=42,
        output_dir=tmp_path / "output",
    )


@pytest.fixture
def store(tmp_path: Path):
    from app.question_gen.run_store import QuestionGenStore

    db_path = tmp_path / "test_qgen.db"
    s = QuestionGenStore(db_path=db_path)
    yield s
    s.close()


# ---------------------------------------------------------------------------
# TestSlotAssignment
# ---------------------------------------------------------------------------


class TestSlotAssignment:
    """_assign_slots 单元测试。"""

    def test_per_type_count(self):
        """验证生成的 slot 总数 = len(task_types) * per_type。"""
        video_ids = ["vid_001", "vid_002"]
        task_types = ["Action Recognition", "Object Recognition", "Counting Problem"]
        per_type = 4

        slots = _assign_slots(video_ids, task_types, per_type)

        assert len(slots) == len(task_types) * per_type

    def test_round_robin_across_videos(self):
        """验证 slot 在视频间轮转分配。"""
        video_ids = ["vid_A", "vid_B"]
        task_types = ["Action Recognition"]
        per_type = 4

        slots = _assign_slots(video_ids, task_types, per_type)

        video_assignments = [s.video_id for s in slots]
        # 应该轮转分配
        assert video_assignments.count("vid_A") == 2
        assert video_assignments.count("vid_B") == 2

    def test_deterministic(self):
        """相同参数产出相同 slot 序列（无随机性）。"""
        video_ids = ["vid_001", "vid_002"]
        task_types = ["Action Recognition", "Temporal Reasoning"]
        per_type = 3

        slots1 = _assign_slots(video_ids, task_types, per_type)
        slots2 = _assign_slots(video_ids, task_types, per_type)

        for s1, s2 in zip(slots1, slots2, strict=True):
            assert s1.slot_id == s2.slot_id
            assert s1.video_id == s2.video_id
            assert s1.task_type == s2.task_type


# ---------------------------------------------------------------------------
# TestProcessOneSlot
# ---------------------------------------------------------------------------


class TestProcessOneSlot:
    """_process_one_slot 集成测试。"""

    @pytest.mark.asyncio
    async def test_happy_path(self, tree, default_config, store, tmp_path):
        """首次生成即通过门控 → 返回 GeneratedQuestion。"""
        vlm = MockVLM()
        llm = MockLLM([_make_gate_pass_response()] * 4)
        sem = asyncio.Semaphore(2)
        used_node_ids: set[str] = set()
        rng = random.Random(42)

        slot = SlotAssignment(
            slot_id="slot_001",
            video_id="vid_L1_000",
            task_type="Action Recognition",
            seq=1,
        )

        run_id = "test-run-001"
        store.record_run_start(run_id, "abc123", "{}")

        result = await _process_one_slot(
            slot=slot,
            tree=tree,
            vlm=vlm,
            llm=llm,
            embed_fn=_mock_embed_fn,
            embed_pool=[],
            store=store,
            config=default_config,
            used_node_ids=used_node_ids,
            rng=rng,
            sem=sem,
            session_id="sess-001",
            run_id=run_id,
        )

        assert result is not None
        assert result.question_id
        # strategy 根据 task_type 决定 skill_target
        from app.question_gen.strategy import get_strategy

        expected_strategy = get_strategy("Action Recognition")
        assert result.skill_target == expected_strategy.skill_target

    @pytest.mark.asyncio
    async def test_retry_on_fail(self, tree, default_config, store, tmp_path):
        """第一次门控失败，第二次通过 → 重出成功。"""
        vlm = MockVLM(
            [
                _make_candidate_json("First question?"),
                _make_candidate_json("Second better question?"),
            ]
        )
        # 第一轮 4 个门有一个 fail，第二轮 4 个门全 pass
        llm_responses = [
            _make_gate_fail_response("answer not grounded"),
            _make_gate_pass_response(),
            _make_gate_pass_response(),
            _make_gate_pass_response(),
            # 第二轮
            _make_gate_pass_response(),
            _make_gate_pass_response(),
            _make_gate_pass_response(),
            _make_gate_pass_response(),
        ]
        llm = MockLLM(llm_responses)
        sem = asyncio.Semaphore(2)
        used_node_ids: set[str] = set()
        rng = random.Random(42)

        slot = SlotAssignment(
            slot_id="slot_002",
            video_id="vid_L1_000",
            task_type="Action Recognition",
            seq=2,
        )

        run_id = "test-run-002"
        store.record_run_start(run_id, "abc123", "{}")

        result = await _process_one_slot(
            slot=slot,
            tree=tree,
            vlm=vlm,
            llm=llm,
            embed_fn=_mock_embed_fn,
            embed_pool=[],
            store=store,
            config=default_config,
            used_node_ids=used_node_ids,
            rng=rng,
            sem=sem,
            session_id="sess-002",
            run_id=run_id,
        )

        assert result is not None
        # 确认第二个问题被接受
        assert "Second" in result.question or result.question_id is not None

    @pytest.mark.asyncio
    async def test_max_retries_none(self, tree, default_config, store, tmp_path):
        """所有重试均失败 → 返回 None。"""
        vlm = MockVLM([_make_candidate_json()] * 5)
        # 所有门控均失败
        llm = MockLLM([_make_gate_fail_response("always fails")] * 20)
        sem = asyncio.Semaphore(2)
        used_node_ids: set[str] = set()
        rng = random.Random(42)

        slot = SlotAssignment(
            slot_id="slot_003",
            video_id="vid_L1_000",
            task_type="Action Recognition",
            seq=3,
        )

        run_id = "test-run-003"
        store.record_run_start(run_id, "abc123", "{}")

        result = await _process_one_slot(
            slot=slot,
            tree=tree,
            vlm=vlm,
            llm=llm,
            embed_fn=_mock_embed_fn,
            embed_pool=[],
            store=store,
            config=default_config,
            used_node_ids=used_node_ids,
            rng=rng,
            sem=sem,
            session_id="sess-003",
            run_id=run_id,
        )

        assert result is None


# ---------------------------------------------------------------------------
# TestPipelineV2
# ---------------------------------------------------------------------------


class TestPipelineV2:
    """run_pipeline_v2 完整流程测试。"""

    @pytest.mark.asyncio
    async def test_full_flow(self, tree, default_config, store, tmp_path):
        """完整管线运行，产出 PipelineResult。"""
        vlm = MockVLM([_make_candidate_json(f"Question {i}?") for i in range(50)])
        llm = MockLLM([_make_gate_pass_response()] * 200)

        config = PipelineConfig(
            per_type=2,
            retry_limit=2,
            heavy_sample_rate=0.5,  # 高比例便于测试
            dedup_threshold=0.85,
            concurrency=2,
            seed=42,
            output_dir=tmp_path / "out",
        )

        # 只用一个 task_type 确保树能满足采样
        result = await run_pipeline_v2(
            video_ids=["vid_L1_000"],
            trees={"vid_L1_000": tree},
            vlm=vlm,
            llm=llm,
            embed_fn=_mock_embed_fn,
            store=store,
            config=config,
            task_types=["Action Recognition"],
        )

        assert isinstance(result, PipelineResult)
        assert result.rejected_count >= 0
        # 至少一题被接受（VLM 和 LLM 全部正常返回）
        assert len(result.accepted) > 0

    @pytest.mark.asyncio
    async def test_progress_resume(self, tree, default_config, store, tmp_path):
        """传入 progress dict → 已处理 slot 被跳过。"""
        llm = MockLLM([_make_gate_pass_response()] * 100)

        config = PipelineConfig(
            per_type=2,
            retry_limit=2,
            heavy_sample_rate=0.0,  # 不做 heavy check
            dedup_threshold=0.85,
            concurrency=2,
            seed=42,
            output_dir=tmp_path / "out",
        )

        # _assign_slots 现在是确定性的（无随机性）
        preview_slots = _assign_slots(
            ["vid_L1_000"],
            ["Action Recognition"],
            config.per_type,
        )
        # 构造 progress 标记所有 slot 已完成
        progress = {s.slot_id: "accepted" for s in preview_slots}

        # 用 progress 跑管线 → 所有 slot 被跳过
        vlm2 = MockVLM([_make_candidate_json()] * 20)

        result2 = await run_pipeline_v2(
            video_ids=["vid_L1_000"],
            trees={"vid_L1_000": tree},
            vlm=vlm2,
            llm=llm,
            embed_fn=_mock_embed_fn,
            store=store,
            config=config,
            task_types=["Action Recognition"],
            progress=progress,
        )

        # 所有 slot 在 progress 中 → VLM 零调用
        assert vlm2._gen_count == 0
        assert len(result2.accepted) == 0

    @pytest.mark.asyncio
    async def test_heavy_check_samples(self, tree, default_config, store, tmp_path):
        """heavy_sample_rate > 0 时有题被抽检。"""
        # LLM 响应：前面是 gate pass，后面增加 heavy check 的步骤响应
        heavy_response = json.dumps(
            {
                "steps": [
                    {"thought": "step 1"},
                    {"thought": "step 2"},
                    {"thought": "step 3"},
                ],
                "answer": "A",
            }
        )
        llm = MockLLM([_make_gate_pass_response()] * 100 + [heavy_response] * 20)
        vlm = MockVLM([_make_candidate_json(f"Q{i}?") for i in range(20)])

        config = PipelineConfig(
            per_type=2,
            retry_limit=2,
            heavy_sample_rate=1.0,  # 100% 抽检
            dedup_threshold=0.85,
            concurrency=2,
            seed=42,
            output_dir=tmp_path / "out",
        )

        result = await run_pipeline_v2(
            video_ids=["vid_L1_000"],
            trees={"vid_L1_000": tree},
            vlm=vlm,
            llm=llm,
            embed_fn=_mock_embed_fn,
            store=store,
            config=config,
            task_types=["Action Recognition"],
        )

        # 100% 抽检 → heavy_sampled 等于 accepted 数
        if result.accepted:
            assert len(result.heavy_sampled) == len(result.accepted)

    @pytest.mark.asyncio
    async def test_store_records_all(self, tree, default_config, store, tmp_path):
        """验证 store 中记录了每道题的生成与门控结果。"""
        vlm = MockVLM([_make_candidate_json()] * 10)
        llm = MockLLM([_make_gate_pass_response()] * 50)

        config = PipelineConfig(
            per_type=2,
            retry_limit=2,
            heavy_sample_rate=0.0,
            dedup_threshold=0.85,
            concurrency=1,
            seed=42,
            output_dir=tmp_path / "out",
        )

        result = await run_pipeline_v2(
            video_ids=["vid_L1_000"],
            trees={"vid_L1_000": tree},
            vlm=vlm,
            llm=llm,
            embed_fn=_mock_embed_fn,
            store=store,
            config=config,
            task_types=["Action Recognition"],
        )

        # 查询 store 中的 items
        cursor = store._conn.execute("SELECT COUNT(*) FROM question_gen_items")
        item_count = cursor.fetchone()[0]
        # 至少有 accepted 数量的记录
        assert item_count >= len(result.accepted)