tests/test_identify.py

"""Unit tests for the NCBI-BLAST identify flow and its Swiss-Prot structure
fallback.

The bug this guards against: ``nr`` is crowded with redundant GenBank / TPA
records (e.g. thousands of Cas9 orthologs as ``tpg|HER…``), none of which carry
a UniProt accession — so a heavily-sequenced protein came back from ``nr`` with
no UniProt hit and the AlphaFold embed dead-ended ("no model available") even
though one plainly exists. ``_attach_structure_homolog`` resolves the closest
reviewed (Swiss-Prot) protein, which always has an AlphaFold model.

Network is fully mocked — no live NCBI calls.
"""
import sys
import types

import pytest

from dee import server


# ── _extract_uniprot: only SwissProt / TrEMBL carry UniProt accessions ─────
class _Aln:
    def __init__(self, hit_id, hit_def="X [Homo sapiens]", hsps=None, length=300):
        self.hit_id = hit_id
        self.hit_def = hit_def
        self.hsps = hsps or []
        self.length = length
        self.accession = hit_id.split("|")[1] if "|" in hit_id else hit_id


class _Hsp:
    def __init__(self, identities=250, align_length=300, expect=1e-40, bits=500):
        self.identities = identities
        self.align_length = align_length
        self.expect = expect
        self.bits = bits


def test_extract_uniprot_swissprot():
    assert server._extract_uniprot(_Aln("sp|Q99ZW2|CAS9_STRP1")) == "Q99ZW2"


def test_extract_uniprot_trembl():
    assert server._extract_uniprot(_Aln("tr|A0A1B2C3D4|A0A1B2C3D4_ECOLI")) == "A0A1B2C3D4"


def test_extract_uniprot_rejects_genbank_refseq_tpa():
    # The exact shapes that buried the Cas9 UniProt entry in the bug report.
    assert server._extract_uniprot(_Aln("tpg|HER4567195.1|")) is None
    assert server._extract_uniprot(_Aln("ref|NP_001234.1|")) is None
    assert server._extract_uniprot(_Aln("gb|AAA12345.1|")) is None
    assert server._extract_uniprot(_Aln("pdb|1ABC|A")) is None


# ── _attach_structure_homolog: fallback only when nr gave no UniProt ───────
def _job(hits):
    j = server.BlastJob(job_id="t", seq_hash="h")
    j.hits = hits
    return j


def test_fallback_pins_uniprot_onto_top_hit(monkeypatch):
    # All nr hits look like the Cas9 case: no UniProt accession anywhere.
    job = _job([
        {"accession": "HER4567195.1", "uniprot": None, "organism": "Streptococcus pyogenes",
         "description": "type II CRISPR RNA-guided endonuclease Cas9"},
        {"accession": "HER4575708.1", "uniprot": None, "organism": "Streptococcus pyogenes",
         "description": "type II CRISPR RNA-guided endonuclease Cas9"},
    ])
    monkeypatch.setattr(server, "_swissprot_uniprot", lambda p: {
        "uniprot": "Q99ZW2", "identity_pct": 99.8,
        "organism": "Streptococcus pyogenes serotype M1", "description": "CRISPR-associated endonuclease Cas9/Csn1",
    })
    server._attach_structure_homolog(job, "MDKKYSIGLD" * 50)
    top = job.hits[0]
    assert top["uniprot"] == "Q99ZW2"
    assert "AF-Q99ZW2-F1" in top["alphafold_url"]
    assert top["alphafold_page"].endswith("/entry/Q99ZW2")
    assert top["uniprot_via"] == "swissprot_homolog"
    assert top["uniprot_identity_pct"] == 99.8
    assert "serotype M1" in top["uniprot_organism"]


def test_fallback_skipped_when_uniprot_already_present(monkeypatch):
    job = _job([{"accession": "X", "uniprot": "P12345"}])
    called = {"n": 0}
    monkeypatch.setattr(server, "_swissprot_uniprot",
                        lambda p: (called.__setitem__("n", called["n"] + 1), {"uniprot": "Q99ZW2"})[1])
    server._attach_structure_homolog(job, "MMM")
    assert called["n"] == 0                         # never ran the fallback
    assert job.hits[0]["uniprot"] == "P12345"       # untouched


def test_fallback_noop_on_empty_hits(monkeypatch):
    job = _job([])
    monkeypatch.setattr(server, "_swissprot_uniprot", lambda p: {"uniprot": "Q99ZW2"})
    server._attach_structure_homolog(job, "MMM")    # must not raise
    assert job.hits == []


def test_fallback_leaves_hits_clean_when_no_homolog(monkeypatch):
    job = _job([{"accession": "X", "uniprot": None}])
    monkeypatch.setattr(server, "_swissprot_uniprot", lambda p: None)
    server._attach_structure_homolog(job, "MMM")
    assert job.hits[0].get("uniprot") is None
    assert "alphafold_url" not in job.hits[0]


# ── _swissprot_uniprot: parses a mocked Swiss-Prot BLAST result ────────────
def test_swissprot_uniprot_parses_mocked_blast(monkeypatch):
    aln = _Aln("sp|Q99ZW2|CAS9_STRP1",
               hit_def="CRISPR-associated endonuclease Cas9/Csn1 [Streptococcus pyogenes]",
               hsps=[_Hsp(identities=1360, align_length=1368)], length=1368)
    record = types.SimpleNamespace(alignments=[aln])

    fake_ncbiww = types.SimpleNamespace(qblast=lambda **kw: "HANDLE")
    fake_ncbixml = types.SimpleNamespace(parse=lambda handle: iter([record]))
    fake_blast = types.ModuleType("Bio.Blast")
    fake_blast.NCBIWWW = fake_ncbiww
    fake_blast.NCBIXML = fake_ncbixml
    monkeypatch.setitem(sys.modules, "Bio.Blast", fake_blast)

    out = server._swissprot_uniprot("MDKKYSIGLD" * 50)
    assert out["uniprot"] == "Q99ZW2"
    assert out["identity_pct"] == pytest.approx(99.4, abs=0.2)
    assert out["organism"] == "Streptococcus pyogenes"