tests/test_dms_seed.py

"""Tests for seeding the commons from public DMS (dee.core.dms_seed).

ESM-free: pure aggregation. We construct synthetic 'assays' where a
substitution's effect is controlled across studies, and check the k-anonymity
floor (≥ MIN_USERS independent assays) and the de-identified output shape.
"""
import pytest

from dee.core import aggregate as _agg
from dee.core.dms_seed import parse_proteingym_csv, seed_rows


def _assay(aid, effects):
    """Build one assay's [(label, score)] where each single-site 'X{i}Y' gets a
    base score = its intended effect (plus a couple of WT-ish low rows so the
    within-assay z-score has spread)."""
    recs = [(lab, val) for lab, val in effects]
    return (aid, recs)


def test_seed_keeps_substitution_seen_in_enough_assays():
    # 'W>L' favorable, measured in 3 independent assays → survives k-anon (=3).
    assays = []
    for i in range(_agg.MIN_USERS):
        assays.append(_assay(f"study{i}", [
            (f"W{10 + i}L", 2.0),     # the favorable W>L
            (f"K{20 + i}D", -1.0),    # spread so z-score is defined
            (f"A{30 + i}G", 0.0),
        ]))
    rows = seed_rows(assays, enforce_gate=False)
    subs = {r["substitution"] for r in rows}
    assert "W>L" in subs
    wl = next(r for r in rows if r["substitution"] == "W>L")
    assert wl["n_users"] >= _agg.MIN_USERS     # backed by ≥ 3 studies
    assert wl["mean_effect"] > 0               # favorable, correct sign


def test_seed_drops_substitution_below_k_anon():
    # 'C>Y' appears in only ONE assay → dropped (privacy floor).
    assays = [
        _assay("only_study", [("C5Y", 3.0), ("K6D", -1.0), ("A7G", 0.0)]),
        _assay("study2", [("K6D", -1.0), ("A7G", 0.5), ("M8I", 0.2)]),
        _assay("study3", [("K6D", -0.8), ("A7G", 0.3), ("M8I", 0.1)]),
    ]
    rows = seed_rows(assays, enforce_gate=False)
    assert "C>Y" not in {r["substitution"] for r in rows}


def test_seed_handles_multi_mutant_colon_labels():
    # ProteinGym multi-mutants ('A1C:D5E') must not crash and should contribute.
    assays = [
        _assay(f"s{i}", [("A1C:D5E", 1.5), ("K2D", -1.0), ("M3I", 0.0)])
        for i in range(_agg.MIN_USERS)
    ]
    rows = seed_rows(assays, enforce_gate=False)
    assert isinstance(rows, list)   # decomposed by the aggregate ridge, no error


def test_seed_empty_when_no_assays():
    assert seed_rows([], enforce_gate=False) == []


def test_parse_proteingym_csv():
    csv_text = "mutant,DMS_score\nA1C,1.2\nD5E,-0.4\nbad,notanumber\n"
    recs = parse_proteingym_csv(csv_text)
    assert recs == [("A1C", 1.2), ("D5E", -0.4)]     # bad row skipped


def test_parse_proteingym_csv_alt_columns():
    csv_text = "mutation,fitness\nW10L,0.9\n"
    assert parse_proteingym_csv(csv_text) == [("W10L", 0.9)]


def test_parse_proteingym_csv_missing_columns():
    assert parse_proteingym_csv("foo,bar\n1,2\n") == []