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

"""单元臂执行任务测试:缓存命中/新鲜跑/INFRA/冻结跳过/题槽并发上限。"""

from __future__ import annotations

import asyncio
from pathlib import Path

import pytest

from app.harness.gate_ladder import BaselineCache
from app.harness.validate import (
    GateSpec,
    _GateRun,
    _QuestionSlots,
    _run_unit_arm,
)
from tests.unit.test_gate_prefix import _PARAMS, _mk_unit  # 复用 fixture


class _FakeLog:
    """假 HarnessLog:query 返回预置 predictions 行。"""

    def __init__(self) -> None:
        self.rows: list[dict] = []

    def query(self, sql: str, params: tuple = ()) -> list[dict]:
        """按 run_id（params[0]）过滤预置行,模拟只读 SELECT。"""
        run_id = params[0]
        return [r for r in self.rows if r["run_id"] == run_id]


def _mk_gate_run(n: int, tmp_path: Path) -> tuple[_GateRun, BaselineCache]:
    """构造 n 个 single 单元的 gate 运行时状态与空基线缓存。"""
    spec = GateSpec(
        task_type="Action Reasoning",
        target_file="action-reasoning.md",
        candidate_content="cand",
        base_skill_content="base",
        units=tuple(_mk_unit(f"q{i}") for i in range(n)),
        gate_run_prefix="r_e1_s0_gate_action-reasoning",
    )
    return _GateRun.from_spec(spec), BaselineCache(tmp_path / "bc.json")


def _fake_run_inference(log: _FakeLog, correct: bool, stop_reason: str = "finished"):
    """构造假推理:把每题结果写进 _FakeLog 并返回带 total 的结果对象。"""

    class _R:
        def __init__(self, run_id: str, total: int) -> None:
            self.run_id = run_id
            self.total = total

    async def _run(questions, *, run_id: str, skills_dir: Path):
        for q in questions:
            log.rows.append(
                {
                    "run_id": run_id,
                    "question_id": q.question_id,
                    "prediction": "A" if correct else "B",
                    "answer": "A",
                    "stop_reason": stop_reason,
                    "steps_json": "[]",
                }
            )
        return _R(run_id, len(questions))

    return _run


@pytest.mark.asyncio
async def test_base_arm_cache_hit_skips_inference(tmp_path) -> None:
    """base 臂缓存命中:不调推理,slot.base 直接就位,infra_denom 不增。"""
    run, cache = _mk_gate_run(1, tmp_path)
    cache.put("Action Reasoning", run.s_hash, "v1", "q0", True)
    called = {"n": 0}

    async def _boom(questions, *, run_id, skills_dir):
        called["n"] += 1
        raise AssertionError("缓存命中不应触发推理")

    slots = _QuestionSlots(4)
    await _run_unit_arm(
        run,
        0,
        "base",
        slots,
        _boom,
        _FakeLog(),
        cache,
        "v1",
        Path("/nonexistent"),
        Path("/nonexistent"),
        _PARAMS,
        0.10,
    )
    assert called["n"] == 0 and run.slots[0].base is True and run.infra_denom == 0


@pytest.mark.asyncio
async def test_base_arm_fresh_run_writes_cache(tmp_path) -> None:
    """base 臂 miss 新鲜跑:结果折叠入 slot 并回写缓存。"""
    run, cache = _mk_gate_run(1, tmp_path)
    log = _FakeLog()
    slots = _QuestionSlots(4)
    await _run_unit_arm(
        run,
        0,
        "base",
        slots,
        _fake_run_inference(log, correct=True),
        log,
        cache,
        "v1",
        tmp_path,
        tmp_path,
        _PARAMS,
        0.10,
    )
    assert run.slots[0].base is True
    assert cache.get("Action Reasoning", run.s_hash, "v1", "q0") is True
    assert run.infra_denom == 1


@pytest.mark.asyncio
async def test_infra_arm_marks_excluded_and_no_cache(tmp_path) -> None:
    """INFRA 臂:标记 infra、errors+1、不写缓存。"""
    run, cache = _mk_gate_run(1, tmp_path)
    log = _FakeLog()
    slots = _QuestionSlots(4)
    await _run_unit_arm(
        run,
        0,
        "base",
        slots,
        _fake_run_inference(log, correct=False, stop_reason="error"),
        log,
        cache,
        "v1",
        tmp_path,
        tmp_path,
        _PARAMS,
        0.10,
    )
    assert run.slots[0].base_infra and run.errors == 1
    assert cache.get("Action Reasoning", run.s_hash, "v1", "q0") is None


@pytest.mark.asyncio
async def test_frozen_run_skips_launch(tmp_path) -> None:
    """已冻结题型的排队臂:直接返回,不占槽不推理。"""
    run, cache = _mk_gate_run(1, tmp_path)
    run.frozen = True
    called = {"n": 0}

    async def _boom(questions, *, run_id, skills_dir):
        called["n"] += 1

    await _run_unit_arm(
        run,
        0,
        "cand",
        _QuestionSlots(4),
        _boom,
        _FakeLog(),
        cache,
        "v1",
        tmp_path,
        tmp_path,
        _PARAMS,
        0.10,
    )
    assert called["n"] == 0


@pytest.mark.asyncio
async def test_question_slots_caps_inflight() -> None:
    """题槽闸:峰值在飞数严格 ≤ 宽度(多槽获取不交错死锁)。"""
    slots = _QuestionSlots(2)
    peak = {"cur": 0, "max": 0}

    async def _job(n: int) -> None:
        await slots.acquire(n)
        peak["cur"] += n
        peak["max"] = max(peak["max"], peak["cur"])
        await asyncio.sleep(0.01)
        peak["cur"] -= n
        slots.release(n)

    await asyncio.gather(*[_job(1) for _ in range(6)], *[_job(2) for _ in range(3)])
    assert peak["max"] <= 2


@pytest.mark.asyncio
async def test_question_slots_rejects_oversized_request() -> None:
    """申请槽数超宽度:fail-fast ValueError 而非自死锁(Codex C2 回归锁)。"""
    slots = _QuestionSlots(1)
    with pytest.raises(ValueError, match="自死锁"):
        await slots.acquire(2)