tests/test_admin_benchmarks_seed.py

"""Tests for the admin maintenance routes that produce the REAL glass-box
receipts server-side (/api/admin/run-benchmarks, /api/admin/seed-commons) —
run in-Space, where the live ESM-2 model and Supabase credentials already
exist. Both are admin-token gated; both are exercised against the ACTUAL
bundled fixture (dee/data/dms_fixtures/) so a malformed manifest/CSV would
fail here too, not just in production.
"""
import json
import re
import types
from pathlib import Path

import pandas as pd
import pytest

from dee import server
from dee.core.dms_seed import parse_proteingym_csv

ADMIN_TOKEN = "test-admin-token-xyz"


@pytest.fixture
def client():
    app = server.create_app()
    app.config.update(TESTING=True)
    return app.test_client()


def _with_admin(monkeypatch):
    monkeypatch.setenv("TURINGDNA_ADMIN_TOKEN", ADMIN_TOKEN)


def _fixtures_dir():
    return Path(server.__file__).resolve().parent / "data" / "dms_fixtures"


# --------------------------------------------------------------------------- #
# auth gate — shared shape for both routes
# --------------------------------------------------------------------------- #
def test_run_benchmarks_requires_admin_token(client, monkeypatch):
    _with_admin(monkeypatch)
    r = client.post("/api/admin/run-benchmarks")
    assert r.status_code == 403
    r2 = client.post("/api/admin/run-benchmarks", headers={"X-Admin-Token": "wrong"})
    assert r2.status_code == 403


def test_seed_commons_requires_admin_token(client, monkeypatch):
    _with_admin(monkeypatch)
    r = client.post("/api/admin/seed-commons")
    assert r.status_code == 403


# --------------------------------------------------------------------------- #
# /api/admin/run-benchmarks — real bundled fixture, fake (fast) scorer
# --------------------------------------------------------------------------- #
def test_run_benchmarks_against_real_fixture(client, monkeypatch, tmp_path):
    _with_admin(monkeypatch)
    real_fixtures = _fixtures_dir()   # resolve BEFORE monkeypatching server.__file__ below
    manifest = json.loads((real_fixtures / "manifest.json").read_text(encoding="utf-8"))
    assert len(manifest) >= 3, "expected the 3 bundled + validated DMS assays"

    # Build a per-sequence fake scorer covering EVERY bundled assay's own real
    # labels (synthetic-but-distinct values, index-based) — a genuine smoke
    # test of all 3 production fixtures end to end (real manifest -> real CSV
    # -> predict_additive -> evaluate_dataset -> summarize -> written file),
    # without needing ESM.
    rx = re.compile(r"^([A-Za-z])(\d+)([A-Za-z*])$")
    df_by_seq = {}
    for a in manifest:
        recs = parse_proteingym_csv((real_fixtures / a["csv"]).read_text(encoding="utf-8"))
        assert len(recs) > 50, f"{a['name']} fixture looks too small to be real"
        rows, seen = [], set()
        for i, (lab, _v) in enumerate(recs):
            m = rx.match(lab)
            key = (int(m.group(2)) - 1, m.group(3).upper())
            if key in seen:
                continue
            seen.add(key)
            rows.append({"position": key[0], "wt_aa": m.group(1).upper(),
                         "mut_aa": key[1], "delta_ll": float(i % 37) * 0.1})
        df_by_seq[a["sequence"]] = pd.DataFrame(rows)

    monkeypatch.setattr(server._scoring, "get_scorer", lambda *a, **kw: "the-scorer")
    monkeypatch.setattr(server._scoring, "score_guarded",
                        lambda scorer, seq: df_by_seq[seq])
    # Write the output somewhere disposable so this test never touches the
    # committed dee/data/benchmarks.json.
    monkeypatch.setattr(server, "__file__",
                        str(tmp_path / "server.py"))  # relocates Path(__file__).parent
    (tmp_path / "data").mkdir()
    import shutil
    shutil.copytree(real_fixtures, tmp_path / "data" / "dms_fixtures")

    r = client.post("/api/admin/run-benchmarks",
                    headers={"X-Admin-Token": ADMIN_TOKEN}, json={"model": "small"})
    assert r.status_code == 200
    body = r.get_json()
    assert body["ok"] is True
    assert body["model"] == "small"
    assert body["generated_at"] is not None
    assert body["summary"]["n_datasets"] == len(manifest)
    names = {d["name"] for d in body["datasets"]}
    assert names == {a["name"] for a in manifest}
    for d in body["datasets"]:
        assert d["n"] > 0

    written = json.loads((tmp_path / "data" / "benchmarks.json").read_text(encoding="utf-8"))
    assert written["summary"]["n_datasets"] == len(manifest)


def test_run_benchmarks_missing_fixtures_404(client, monkeypatch, tmp_path):
    _with_admin(monkeypatch)
    monkeypatch.setattr(server, "__file__", str(tmp_path / "server.py"))
    (tmp_path / "data").mkdir()   # no dms_fixtures subdir
    r = client.post("/api/admin/run-benchmarks", headers={"X-Admin-Token": ADMIN_TOKEN})
    assert r.status_code == 404


def test_run_benchmarks_surfaces_per_assay_failure(client, monkeypatch, tmp_path):
    """One assay's scoring blows up (e.g. a real ESM error on the Space) — the
    other two must still complete, and the failure must be visible in the
    response itself (name + error), not just swallowed into a shorter
    'datasets' list with no explanation. This is what a real diagnosis of
    'why did only 1 of 3 datasets come back' should read directly off the
    curl output instead of requiring server-log access."""
    _with_admin(monkeypatch)
    real_fixtures = _fixtures_dir()
    manifest = json.loads((real_fixtures / "manifest.json").read_text(encoding="utf-8"))
    assert len(manifest) >= 3

    def flaky_score_guarded(scorer, seq):
        if len(seq) > 500:   # the two longer real fixtures (PABP 577aa, DLG4 724aa)
            raise RuntimeError("boom: simulated real scoring failure")
        return pd.DataFrame([{"position": 0, "wt_aa": seq[0], "mut_aa": "Z", "delta_ll": 0.1}])

    monkeypatch.setattr(server._scoring, "get_scorer", lambda *a, **kw: "the-scorer")
    monkeypatch.setattr(server._scoring, "score_guarded", flaky_score_guarded)
    monkeypatch.setattr(server, "__file__", str(tmp_path / "server.py"))
    (tmp_path / "data").mkdir()
    import shutil
    shutil.copytree(real_fixtures, tmp_path / "data" / "dms_fixtures")

    r = client.post("/api/admin/run-benchmarks",
                    headers={"X-Admin-Token": ADMIN_TOKEN}, json={"model": "small"})
    assert r.status_code == 200
    body = r.get_json()
    long_assays = [a["name"] for a in manifest if len(a["sequence"]) > 500]
    assert len(long_assays) == 2, "expected exactly PABP + DLG4 to be >500aa"
    assert body["summary"]["n_datasets"] == len(manifest) - len(long_assays)
    assert "failed" in body and len(body["failed"]) == len(long_assays)
    failed_names = {f["name"] for f in body["failed"]}
    assert failed_names == set(long_assays)
    for f in body["failed"]:
        assert "boom: simulated real scoring failure" in f["error"]


# --------------------------------------------------------------------------- #
# /api/admin/seed-commons — real bundled fixture, mocked Supabase write
# --------------------------------------------------------------------------- #
def test_seed_commons_against_real_fixture(client, monkeypatch):
    _with_admin(monkeypatch)
    captured = {}

    def fake_replace(rows, source="user"):
        captured["rows"] = rows
        return {"ok": True}

    monkeypatch.setattr(server._auth, "replace_mutation_priors", fake_replace)
    r = client.post("/api/admin/seed-commons", headers={"X-Admin-Token": ADMIN_TOKEN})
    assert r.status_code == 200
    body = r.get_json()
    assert body["ok"] is True
    assert body["contributing_assays"] == 3
    # 3 independent assays clears the k-anonymity floor (MIN_USERS=3): every
    # single-mutant scan covers close to all 19 substitutions at every
    # position, so most of the 380 possible substitution TYPES show up in
    # all three proteins' data and survive. Proves the privacy floor is
    # exactly a floor, not a permanent block, once there's real independent
    # coverage — and that it produces real, usable rows on real data.
    assert body["substitutions"] > 0
    assert len(captured["rows"]) == body["substitutions"]
    for row in captured["rows"]:
        assert row["n_users"] >= 3   # every kept row backed by all 3 independent assays
        assert ">" in row["substitution"]


def test_seed_commons_survives_one_bad_csv(client, monkeypatch, tmp_path):
    """A single unreadable/corrupt fixture CSV must not 500 the whole request
    (there was previously no try/except around this loop at all) — it should
    degrade to the other assays and report the failure by name."""
    _with_admin(monkeypatch)
    real_fixtures = _fixtures_dir()
    monkeypatch.setattr(server._auth, "replace_mutation_priors", lambda rows, source="user": {"ok": True})
    monkeypatch.setattr(server, "__file__", str(tmp_path / "server.py"))
    (tmp_path / "data").mkdir()
    import shutil
    dest = tmp_path / "data" / "dms_fixtures"
    shutil.copytree(real_fixtures, dest)
    manifest = json.loads((dest / "manifest.json").read_text(encoding="utf-8"))
    broken_name = manifest[0]["name"]
    (dest / manifest[0]["csv"]).unlink()   # simulate a missing/corrupt fixture file

    r = client.post("/api/admin/seed-commons", headers={"X-Admin-Token": ADMIN_TOKEN})
    assert r.status_code == 200
    body = r.get_json()
    assert body["contributing_assays"] == len(manifest) - 1
    assert "failed" in body and len(body["failed"]) == 1
    assert body["failed"][0]["name"] == broken_name


def test_seed_commons_gated_before_effective_date(client, monkeypatch):
    _with_admin(monkeypatch)
    import datetime as _dt
    from dee.core import aggregate as _agg
    monkeypatch.setattr(_agg, "EFFECTIVE_DATE", _dt.date(2099, 1, 1))
    r = client.post("/api/admin/seed-commons", headers={"X-Admin-Token": ADMIN_TOKEN})
    assert r.status_code == 423
    assert r.get_json()["error"] == "gated"