tests/test_base_editor.py

"""Base-editing design (dee/core/base_editor.py).

Pins the editor library, sequence-only edit prediction (incl. bystander
detection), and — most importantly — the frame-aware codon consequence,
which is what tells a customer "this guide installs a premature stop".
A wrong AA call here would be as damaging as a wrong off-target, so the
codon math is tested against hand-verified examples and is required to
return None rather than guess when the alignment doesn't check out.
"""
import pytest

from dee.core import base_editor as BE


# ─────────────────────────── editor library ───────────────────────────

def test_library_shape():
    be4 = BE.get_base_editor("be4max")
    assert be4.kind == "CBE" and be4.target_base == "C" and be4.result_base == "T"
    assert be4.window == (4, 8)
    abe = BE.get_base_editor("abe7.10")
    assert abe.kind == "ABE" and abe.target_base == "A" and abe.result_base == "G"
    assert BE.get_base_editor("nope") is None


def test_list_filters_by_kind():
    assert {e.id for e in BE.list_base_editors("CBE")} == {"be4max", "be3", "evocda_be4max"}
    assert {e.id for e in BE.list_base_editors("ABE")} == {"abe8e", "abe7.10"}


def test_activity_profiles_bounded_and_peaked():
    for ed in BE.list_base_editors():
        lo, hi = ed.window
        assert set(ed.activity) == set(range(lo, hi + 1))
        assert all(0.0 <= v <= 1.0 for v in ed.activity.values())
        assert max(ed.activity.values()) == pytest.approx(1.0)


# ─────────────────────────── edit prediction ───────────────────────────

def test_single_cbe_edit_in_window():
    spacer = "AAAC" + "A" * 16          # C at position 4 only
    p = BE.predict_base_edits(spacer, "be4max")
    assert p.n_targets == 1
    assert p.edits[0][:3] == (4, "C", "T")
    assert p.editability == pytest.approx(0.35)   # window edge (pos 4)
    assert p.has_bystander is False
    assert p.summary == "C4→T"


def test_bystander_flagged_and_editability_is_strongest():
    spacer = "AAACA" + "C" + "A" * 14    # C at positions 4 and 6
    p = BE.predict_base_edits(spacer, "be4max")
    assert p.n_targets == 2
    assert p.has_bystander is True
    assert p.editability == pytest.approx(1.0)    # pos 6 is the window peak
    assert "bystander" in p.summary


def test_no_target_in_window():
    p = BE.predict_base_edits("A" * 20, "be4max")
    assert p.n_targets == 0 and p.editability == 0.0
    assert "no C" in p.summary


def test_abe_targets_adenine():
    spacer = "GGGA" + "G" * 16           # A at position 4
    p = BE.predict_base_edits(spacer, "abe7.10")
    assert p.n_targets == 1
    assert p.edits[0][:3] == (4, "A", "G")


def test_unknown_editor_returns_empty():
    p = BE.predict_base_edits("ACGT" * 5, "not_an_editor")
    assert p.n_targets == 0 and p.edits == []


# ─────────────────────────── codon consequence ───────────────────────────

def test_istop_cga_to_tga():
    # CGA (Arg) --C→T at codon position 1--> TGA (stop). The CRISPR-STOP move.
    c = BE.codon_consequence("CGA" + "GGG", 0, "C", "T")
    assert c.creates_stop is True
    assert c.kind == "nonsense"
    assert (c.aa_before, c.aa_after) == ("R", "*")
    assert c.label == "R1* (stop gained)"


@pytest.mark.parametrize("codon,aa", [("CAA", "Q"), ("CAG", "Q"), ("CGA", "R")])
def test_all_cbe_istop_codons(codon, aa):
    # The three codons CBE can turn into a stop by a single C→T.
    c = BE.codon_consequence(codon + "TTT", 0, "C", "T")
    assert c.creates_stop and c.aa_before == aa and c.aa_after == "*"


def test_silent_edit():
    # GCT (Ala) --T→C at position 3--> GCC (still Ala).
    c = BE.codon_consequence("GCT" + "AAA", 2, "T", "C")
    assert c.kind == "silent" and c.aa_before == c.aa_after == "A"
    assert "silent" in c.label


def test_missense_edit():
    # GCT (Ala) --G→A at position 1--> ACT (Thr).
    c = BE.codon_consequence("GCT" + "AAA", 0, "G", "A")
    assert c.kind == "missense" and (c.aa_before, c.aa_after) == ("A", "T")


def test_stop_loss_edit():
    # TGA (stop) --A→G at position 3--> TGG (Trp).
    c = BE.codon_consequence("TGA" + "GCT", 2, "A", "G")
    assert c.kind == "stop_loss" and c.creates_stop is False
    assert (c.aa_before, c.aa_after) == ("*", "W")


def test_residue_numbering_uses_codon_index():
    # Second codon → residue 2.
    c = BE.codon_consequence("GGG" + "CGA", 3, "C", "T")
    assert c.label == "R2* (stop gained)"


def test_refuses_to_guess_on_mismatch_or_out_of_range():
    # ref_base disagrees with the CDS → None (no fabricated consequence).
    assert BE.codon_consequence("GCT", 0, "A", "T") is None
    # index past the end → None.
    assert BE.codon_consequence("GCT", 99, "T", "C") is None
    # incomplete trailing codon → None.
    assert BE.codon_consequence("GC", 0, "G", "A") is None


# ─────────────── integration: find_guides(mode="base_edit") ───────────────
# Locks the full path — edit prediction + coordinate mapping (both strands)
# + frame auto-detection + codon translation — against a fixed in-frame CDS.

from dee.core import crispr as C  # noqa: E402

_CDS = "".join([
    "ATG", "GCT", "CGA", "GAT", "CAG", "GGC", "CGT", "CAA", "GCC", "TGG",
    "CGC", "GAA", "CGG", "CTG", "CGA", "GCA", "CAG", "GAT", "CGT", "TAA",
])  # 60 nt, starts ATG, ends TAA, no internal stop → auto-detected frame


def test_base_edit_mode_populates_fields_and_ranks_by_editability():
    guides = C.find_guides(_CDS, enzyme="cas9", max_results=200,
                           mode="base_edit", base_editor="be4max")
    assert guides
    for g in guides:
        assert g.be_editor == "be4max"
        assert 0.0 <= g.be_editability <= 1.0
    eds = [g.be_editability for g in guides]
    assert eds == sorted(eds, reverse=True)        # editability-first ranking


def test_base_edit_detects_istop_on_forward_strand():
    guides = C.find_guides(_CDS, enzyme="cas9", max_results=200,
                           mode="base_edit", base_editor="be4max")
    stops = {g.be_aa_change for g in guides if g.be_creates_stop}
    # CAG (Q) → TAG (*) via a C→T in the window, at residues 5 and 8.
    assert "Q5* (stop gained)" in stops
    assert "Q8* (stop gained)" in stops
    for g in guides:
        if g.be_creates_stop:
            assert g.be_aa_change.endswith("* (stop gained)")


def test_base_edit_reverse_strand_aa_mapping():
    # A '-' strand guide editing the coding-strand G→A: validates the
    # reverse-strand coordinate + complement transform end-to-end.
    guides = C.find_guides(_CDS, enzyme="cas9", max_results=200,
                           mode="base_edit", base_editor="be4max")
    changes = {g.be_aa_change for g in guides if g.strand == "-"}
    assert "E12K" in changes


def test_base_edit_flags_bystander():
    guides = C.find_guides(_CDS, enzyme="cas9", max_results=200,
                           mode="base_edit", base_editor="be4max")
    assert any(g.be_has_bystander for g in guides)


def test_no_aa_change_without_frame():
    # A random (non-CDS) sequence → edits predicted, but no AA call.
    import random
    rng = random.Random(11)
    seq = "".join(rng.choice("ACGT") for _ in range(400))
    guides = C.find_guides(seq, enzyme="cas9", max_results=50,
                           mode="base_edit", base_editor="be4max")
    assert guides
    assert all(g.be_aa_change == "" for g in guides)   # frame unknown → no guess