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

"""训练循环诊断注入：_run_diagnosis 按 batch question_ids 加载真实树注入 run_diagnosis。

覆盖诊断 tree_data 断链修复的训练循环侧（Task 4）：验证 runner._run_diagnosis 会把
batch 涉及 video 的真实树作为 tree_data 传给 core.run_diagnosis，而非空 dict。
"""

from __future__ import annotations

import json
from pathlib import Path  # noqa: TC003 — 运行时 tmp_path 标注使用
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from app.harness.config import RunConfig
from app.harness.runner import Runner
from core.evolution.types import DiagnosisResult

# 真实样本：question 604-2 属于 video 0RxMZBLeqRI（111 节点树）。
_REAL_QUESTION_ID = "604-2"
_REAL_VIDEO_ID = "0RxMZBLeqRI"


def _empty_diagnosis_result(run_id: str = "infer_adhoc") -> DiagnosisResult:
    """构造仅含 run_id 的空诊断结果（其余字段走 dataclass 默认值）。"""
    return DiagnosisResult(run_id=run_id)


def _base_config(workspace_dir: Path, store_dir: Path) -> RunConfig:
    """构造 diagnose 模式 RunConfig，所有必填字段给测试默认值。

    参数:
        workspace_dir: 已写入 manifest 的 workspace 根目录。
        store_dir: 真实 store 根目录（含 questions/ 与 videos/）。
    """
    return RunConfig(
        workspace_dir=workspace_dir,
        store_dir=store_dir,
        mode="diagnose",
        concurrency=1,
        max_steps=5,
        skill_mode="none",
        n_samples=0,
        questions="benchmarks/Video-MME",
        skills_version="v1",
        prompts_version="v1",
        epochs=1,
        diag_size=10,
        diag_correct_ratio=0.5,
        val_size=24,
        val_correct_ratio=0.5,
        edit_budget_start=5,
        edit_budget_end=2,
        batch_size=5,
        min_class_per_batch=2,
        eval_min_per_class=2,
        early_stop_patience=3,
        test_size=10,
        use_slow_momentum=False,
        gate_e_confirm=20.0,
        gate_e_provisional=3.0,
        gate_w_net_min=2,
        gate_delta_min=0.02,
        gate_lambda_dir=-0.642,
        gate_e_rollback=10.0,
        gate_block=8,
        gate_n_max=40,
        gate_p_low=0.05,
        gate_p_high=0.95,
        gate_probe_quota=0.2,
        gate_gamma_decay=0.9,
        gate_cooldown_steps=2,
        gate_guard_err=0.10,
        skill_update_mode="patch",
        appendix_consolidate_threshold=6,
        run_id="infer_adhoc",
    )


@pytest.fixture
def runner_with_real_store(tmp_path: Path) -> Runner:
    """构造 diagnose 模式 runner，questions/videos 指向真实 store。

    manifest.store 写真实 store 绝对路径，使 resolve_paths 的 questions_dir
    命中含 604-2 的题库、store_dir 命中 0RxMZBLeqRI 的真实 tree.json。
    """
    store_dir = Path(__file__).resolve().parents[2] / "store"
    ws = tmp_path / "ws"
    ws.mkdir()
    (ws / "skills" / "v1").mkdir(parents=True)
    manifest = {
        "name": "ws",
        "created_at": "",
        "store": store_dir.as_posix(),
        "current": {
            "videos": "videos",
            "questions": "questions/benchmarks/Video-MME",
            "skills": "skills/v1",
            "prompts": "prompts/v1",
        },
        "history": [],
    }
    (ws / "manifest.json").write_text(json.dumps(manifest))

    return Runner(
        _base_config(ws, store_dir),
        llm=MagicMock(),
        evolve_llm=MagicMock(),
        vlm=MagicMock(),
        telemetry=MagicMock(),
    )


@pytest.mark.asyncio
async def test_run_diagnosis_injects_tree_data(runner_with_real_store: Runner) -> None:
    """_run_diagnosis 把 batch question_ids 对应 video 的真实树注入 run_diagnosis。"""
    captured: dict[str, object] = {}

    async def _fake_run_diagnosis(**kwargs: object) -> DiagnosisResult:
        captured["tree_data"] = kwargs["tree_data"]
        return _empty_diagnosis_result()

    with patch(
        "core.evolution.diagnose.run_diagnosis",
        new=AsyncMock(side_effect=_fake_run_diagnosis),
    ):
        await runner_with_real_store._run_diagnosis("infer_adhoc", question_ids=[_REAL_QUESTION_ID])

    tree_data = captured["tree_data"]
    assert _REAL_VIDEO_ID in tree_data
    assert tree_data[_REAL_VIDEO_ID]["nodes"]


def _fake_question(question_id: str, video_id: str) -> object:
    """构造仅设置 question_id/video_id 的 GeneratedQuestion（其余字段给占位默认）。"""
    from core.types import GeneratedQuestion

    return GeneratedQuestion(
        question_id=question_id,
        video_id=video_id,
        task_type="Action Reasoning",
        question="问题",
        options=("A", "B", "C", "D"),
        answer="A",
        source_nodes=(),
        difficulty="medium",
    )


@pytest.mark.asyncio
async def test_run_diagnosis_full_scan_loads_all_video_trees(
    runner_with_real_store: Runner,
) -> None:
    """question_ids=None（全量诊断）路径：加载全部 questions 涉及 video 的树。

    轻量验证：patch load_benchmark 返回两个不同 video 的假题，patch
    load_tree_data_for_videos 捕获 video_ids，断言全量路径覆盖全部 video。
    """
    fake_qs = [
        _fake_question("q-a", "vA"),
        _fake_question("q-b", "vB"),
    ]
    captured: dict[str, object] = {}

    def _cap(store_dir: Path, video_ids: list[str]) -> dict[str, object]:
        captured["video_ids"] = list(video_ids)
        return {v: {"nodes": {}} for v in video_ids}

    with (
        patch("app.question_gen.load_benchmark", return_value=fake_qs),
        patch("app.harness.tree_nodes.load_tree_data_for_videos", side_effect=_cap),
        patch(
            "core.evolution.diagnose.run_diagnosis",
            new=AsyncMock(return_value=_empty_diagnosis_result()),
        ),
    ):
        await runner_with_real_store._run_diagnosis("infer_adhoc", question_ids=None)

    assert set(captured["video_ids"]) == {"vA", "vB"}