ely-eye/bench/tests/test_joint.py

"""JointRunner must rank exactly like ArmRunner, and the rankings cache must resume.

The ablation switched to a single pass that computes each signal once. If that pass
drifted from the per-arm path, the baseline would silently measure a different system.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import pytest

from bench import systems
from bench.run import (
    Query,
    fingerprint,
    latency,
    load_rankings,
    remeasure_latency,
    scale_arm_names,
    scale_probes,
    write_markdown,
)
from bench.systems import ARMS, ArmRunner, Atom, JointRunner
from bench.tests.test_systems import ATOMS

MODEL_ARMS = ["bm25_only", "graph_only", "dense_only", "no_rerank", "fusion_rerank"]


def _queries() -> list[str]:
    return ["alpha report", "beta invoice", "gamma", ""]


@pytest.mark.parametrize("arm", MODEL_ARMS)
def test_joint_ranking_matches_the_per_arm_runner(arm: str, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(systems, "INDEX_DIR", tmp_path)
    queries = _queries()
    joint = JointRunner([arm], ATOMS, systems.MODELS)
    per_arm = ArmRunner(arm, ATOMS, systems.MODELS)
    for query in queries:
        expected = per_arm.rank(query)
        ranked, seconds = joint.rank_all(query)[arm]
        assert ranked == expected, f"{arm} drifted from the per-arm path on {query!r}"
        assert seconds >= 0.0


def test_joint_runner_computes_the_sparse_signal_once_per_query(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    monkeypatch.setattr(systems, "INDEX_DIR", tmp_path)
    joint = JointRunner(["bm25_only", "no_rerank", "fusion_rerank"], ATOMS, systems.MODELS)
    calls = 0
    assert joint._sparse is not None
    original = joint._sparse.scores

    def counting_scores(query: str):
        nonlocal calls
        calls += 1
        return original(query)

    monkeypatch.setattr(joint._sparse, "scores", counting_scores)
    joint.rank_all("alpha report")
    assert calls == 1, "three arms that all need BM25 must still score the corpus once"


def test_joint_runner_rejects_unknown_or_empty_arm_lists() -> None:
    with pytest.raises(SystemExit, match="unknown arm"):
        JointRunner(["vibes"], ATOMS, systems.MODELS)
    with pytest.raises(SystemExit, match="at least one arm"):
        JointRunner([], ATOMS, systems.MODELS)


def test_rankings_cache_resumes_and_refuses_a_mismatched_run(tmp_path: Path) -> None:
    atoms = [Atom(atom_id=f"a{index}", source_id=f"s{index}", text=f"text {index}") for index in range(3)]
    queries = [
        Query(
            query_id=f"q{index}",
            text="t",
            modality="pdf",
            gold=frozenset({"s0"}),
            holder="h",
            answerable=True,
            as_of=None,
        )
        for index in range(2)
    ]
    mark = fingerprint(["bm25_only"], queries, atoms)
    path = tmp_path / "rankings.jsonl"
    assert load_rankings(path, mark) == {}
    rows = [json.dumps({"fingerprint": mark}), json.dumps({"query_id": "q0", "arms": {"bm25_only": {"ranked": ["s0"]}}})]
    path.write_text("\n".join(rows) + "\n", encoding="utf-8")
    assert load_rankings(path, mark) == {"q0": {"bm25_only": {"ranked": ["s0"]}}}
    with pytest.raises(SystemExit, match="different run"):
        load_rankings(path, "other")
    path.write_text("\n".join([rows[0], "{not json"]) + "\n", encoding="utf-8")
    with pytest.raises(SystemExit, match="not JSON"):
        load_rankings(path, mark)


def test_latency_needs_samples() -> None:
    assert latency([0.5, 1.0])["p50"] == pytest.approx(0.75)
    with pytest.raises(SystemExit, match="no latency samples"):
        latency([])


def test_every_arm_declares_the_signals_it_consumes() -> None:
    assert set(systems.ARM_COMPONENTS) == set(ARMS)


def _probe_set(count: int) -> list[Query]:
    return [
        Query(
            query_id=f"q{index}",
            text="t",
            modality="pdf",
            gold=frozenset({f"s{index}"}),
            holder="h",
            answerable=index % 2 == 0,
            as_of=None,
        )
        for index in range(count)
    ]


def test_scale_probes_ignore_gold_so_every_corpus_size_is_measured() -> None:
    queries = _probe_set(100)
    probes = scale_probes(queries, 10)
    assert len(probes) == 10
    assert all(query.answerable for query in probes)
    assert probes == scale_probes(queries, 10), "the curve must be reproducible across runs"
    assert {query.query_id for query in probes} <= {q.query_id for q in queries if q.answerable}
    with pytest.raises(SystemExit, match="at least one probe"):
        scale_probes(queries, 0)
    with pytest.raises(SystemExit, match="needs answerable"):
        scale_probes([], 10)


def _overall() -> dict[str, float]:
    return {
        "recall_at_1": 0.3,
        "recall_at_5": 0.4,
        "recall_at_10": 0.5,
        "ndcg_at_10": 0.42,
        "mrr": 0.4,
        "hit_at_10": 0.52,
        "hit_at_10_wilson_lower": 0.45,
        "abstention_rate_answerable": 0.0,
        "abstention_rate_unanswerable": 0.0,
        "top_hit_from_holder_unanswerable": 0.2,
    }


def _report() -> dict:
    return {
        "meta": {
            "generated_at": "2026-01-01T00:00:00+00:00",
            "git_commit": "a" * 40,
            "corpus_sources": 10,
            "queries": 16,
            "answerable": 8,
            "machine": {
                "system": "Linux",
                "release": "x",
                "cpus": 32,
                "threads_per_query": 4,
                "load_average_1_5_15": [1.0, 2.0, 3.0],
            },
            "models": {"dense": "d", "rerank": "r"},
        },
        "chance": {"recall_at_1": 0.0, "recall_at_5": 0.0, "recall_at_10": 0.0},
        "arms": [
            {
                "arm": "bm25_only",
                "overall": _overall(),
                "by_modality": {"pdf": _overall()},
                "latency_seconds": {"p50": 0.1, "p95": 0.2},
            }
        ],
        "live_product": {
            "queries": 2,
            "top_k": 10,
            "recall_at_10": 0.87,
            "mrr": 0.74,
            "hit_at_10": 0.87,
            "hit_at_10_wilson_lower": 0.82,
            "p50": 1.0,
            "p95": 1.0,
            "mean": 1.0,
            "max": 1.0,
            "top10_overlap_with_bm25_only": 0.0,
            "identical_top10": 0.0,
            "distinct_sources_in_top10": {"mean": 7.5, "min": 1, "max": 10},
        },
        "scale": [
            {"sources": 500, "atoms": 900, "queries": 4, "note": "latency only", "bm25_only": {"p50": 0.001, "p95": 0.002}}
        ],
    }


def test_the_report_writer_survives_a_scale_note(tmp_path: Path) -> None:
    path = tmp_path / "BASELINE.md"
    write_markdown(_report(), path)
    text = path.read_text(encoding="utf-8")
    assert "| 500 | 900 | 0.001/0.002 |" in text
    assert "_latency only_" in text
    assert "load average 1.0 on 32 cores" in text
    assert "## unanswerable queries (50.0% of the suite)" in text
    assert "recall@10 0.870" in text


def test_scale_arm_names_reads_only_latency_dicts() -> None:
    assert scale_arm_names(_report()["scale"]) == ["bm25_only"]
    with pytest.raises(SystemExit, match="no measured arm"):
        scale_arm_names([{"sources": 500, "note": "latency only"}])


def test_remeasure_latency_replaces_cached_seconds_with_a_warm_sample() -> None:
    arms: list[dict[str, Any]] = [{"arm": "bm25_only", "latency_seconds": {"p50": 99.0}}]
    remeasure_latency(arms, JointRunner(["bm25_only"], ATOMS, systems.MODELS), _probe_set(8), 4)
    measured = arms[0]["latency_seconds"]
    assert measured["p50"] < 99.0
    assert measured["queries"] == 4
    assert "warm models" in measured["note"]