tests/test_benchmark.py

"""Tests for the validation harness (dee.core.benchmark) — the receipts engine.

Predictions are injected, so we can assert the exact correlation / precision
behaviour without ESM. These pin the numbers we'd publish as our benchmark.
"""
import numpy as np
import pytest

import pandas as pd

from dee.core.benchmark import (
    DatasetResult,
    evaluate_dataset,
    predict_additive,
    spearman,
    summarize,
    top_decile_precision,
)


def _fake_scores_df():
    # position 0-indexed: pos 0 A->C = 1.0, pos 4 D->E = 0.5
    return pd.DataFrame([
        {"position": 0, "wt_aa": "A", "mut_aa": "C", "delta_ll": 1.0},
        {"position": 4, "wt_aa": "D", "mut_aa": "E", "delta_ll": 0.5},
        {"position": 4, "wt_aa": "D", "mut_aa": "G", "delta_ll": -0.3},
    ])


def test_predict_additive_single_and_multi_site():
    df = _fake_scores_df()
    preds = predict_additive(df, ["A1C", "A1C:D5E", "A1C,D5E"])
    assert preds[0] == pytest.approx(1.0)
    assert preds[1] == pytest.approx(1.5)   # colon-separated multi-site sums
    assert preds[2] == pytest.approx(1.5)   # comma-separated too


def test_predict_additive_nan_for_unscorable_label():
    df = _fake_scores_df()
    preds = predict_additive(df, ["Z9Q", "A1C:Z9Q", "not-a-label"])
    assert all(np.isnan(p) for p in preds)


def test_spearman_perfect_and_anti():
    a = list(range(20))
    assert spearman(a, a) == pytest.approx(1.0)
    assert spearman(a, list(reversed(a))) == pytest.approx(-1.0)


def test_spearman_monotonic_nonlinear_is_one():
    a = list(range(1, 11))
    b = [x ** 3 for x in a]            # monotonic → rank corr = 1 even if nonlinear
    assert spearman(a, b) == pytest.approx(1.0)


def test_spearman_handles_ties():
    a = [1, 1, 2, 2, 3, 3]
    b = [1, 1, 2, 2, 3, 3]
    assert spearman(a, b) == pytest.approx(1.0)


def test_spearman_none_when_degenerate():
    assert spearman([1, 2], [3, 4]) is None          # < 3 points
    assert spearman([5, 5, 5, 5], [1, 2, 3, 4]) is None  # no variance in a


def test_top_decile_precision_perfect_alignment():
    rng = np.random.default_rng(0)
    measured = rng.normal(size=100)
    predicted = measured.copy()                       # perfect ranking
    # top 10% predicted are exactly the top 10% measured, all within top 25%.
    assert top_decile_precision(predicted, measured) == pytest.approx(1.0)


def test_top_decile_precision_anti_alignment_is_low():
    measured = np.linspace(0, 1, 100)
    predicted = -measured                             # worst possible ranking
    assert top_decile_precision(predicted, measured) == pytest.approx(0.0)


def test_top_decile_precision_none_when_too_small():
    assert top_decile_precision([1, 2, 3], [3, 2, 1]) is None


def test_evaluate_dataset_and_summarize():
    rng = np.random.default_rng(1)
    m1 = rng.normal(size=60)
    p1 = m1 + rng.normal(scale=0.3, size=60)          # good but noisy predictor
    m2 = rng.normal(size=40)
    p2 = -m2                                            # a hard/anti assay
    r1 = evaluate_dataset("assayA", "P1", p1, m1, source="doi:1")
    r2 = evaluate_dataset("assayB", "P2", p2, m2, source="doi:2")
    assert r1.spearman > 0.6
    assert r2.spearman < 0
    s = summarize([r1, r2])
    assert s["n_datasets"] == 2
    assert s["n_variants"] == 100
    assert s["median_spearman"] is not None
    # as_dict rounds + is JSON-safe
    d = r1.as_dict()
    assert set(d) == {"name", "protein", "n", "spearman", "top_decile_precision", "source"}


def test_summarize_empty():
    s = summarize([])
    assert s == {"n_datasets": 0, "n_variants": 0,
                 "median_spearman": None, "median_top_decile_precision": None}


def test_benchmarks_route_public_and_shaped():
    from dee import server
    app = server.create_app()
    app.config.update(TESTING=True)
    body = app.test_client().get("/api/benchmarks").get_json()   # no auth — public
    assert body["ok"] is True
    assert "summary" in body and "datasets" in body
    assert set(body["summary"]) == {"n_datasets", "n_variants",
                                    "median_spearman", "median_top_decile_precision"}
    # n_datasets/datasets length must agree, whatever the current bundled state is
    # (empty before the first real run, or real numbers once /api/admin/run-
    # benchmarks has populated it — this route never fabricates either way).
    assert body["summary"]["n_datasets"] == len(body["datasets"])


def test_benchmarks_route_honest_empty_when_file_absent(monkeypatch, tmp_path):
    from dee import server
    monkeypatch.setattr(server, "__file__", str(tmp_path / "server.py"))
    app = server.create_app()
    app.config.update(TESTING=True)
    body = app.test_client().get("/api/benchmarks").get_json()
    assert body["ok"] is True
    assert body["summary"]["n_datasets"] == 0
    assert body["datasets"] == []
    assert body["generated_at"] is None