tests/test_plasmid.py

"""Unit tests for Plasmid Studio core (dee/core/plasmid.py).
Uses Biopython locally (no network)."""
import pytest

from dee.core import plasmid as PL

# T7 promoter (18) + EcoRI site + 20 + EcoRI site + 20  = 70 bp, 2 EcoRI sites.
SEQ = ("TAATACGACTCACTATAG" + "GAATTC" + "AAAAGGGGAAAAGGGGAAAA"
       + "GAATTC" + "CCCCTTTTCCCCTTTTCCCC")


def test_revcomp_gc():
    assert PL.revcomp("ATGC") == "GCAT"
    assert PL.gc_percent("GGCC") == 100.0
    assert PL.gc_percent("ATAT") == 0.0


def test_parse_raw_dna():
    p = PL.parse_plasmid(SEQ)
    assert p["source"] == "raw"
    assert p["length"] == len(SEQ)
    assert p["topology"] == "circular"
    assert p["features"] == []


def test_parse_fasta():
    p = PL.parse_plasmid(">myseq description here\n" + SEQ)
    assert p["source"] == "fasta"
    assert p["name"] == "myseq"
    assert p["sequence"] == SEQ


def test_parse_topology_override():
    p = PL.parse_plasmid(SEQ, topology="linear")
    assert p["topology"] == "linear"


def test_parse_rejects_empty_and_nondna():
    with pytest.raises(ValueError):
        PL.parse_plasmid("   ")
    with pytest.raises(ValueError):
        PL.parse_plasmid("hello world this is prose not dna @@@")


def test_parse_rejects_oversize(monkeypatch):
    monkeypatch.setattr(PL, "MAX_LEN", 50)
    with pytest.raises(ValueError):
        PL.parse_plasmid("A" * 60)


def test_parse_rejects_oversize_input_text(monkeypatch):
    monkeypatch.setattr(PL, "MAX_INPUT_TEXT", 1000)
    with pytest.raises(ValueError):
        PL.parse_plasmid("A" * 2000)


# ── sanitize_features (audit hardening) ─────────────────────────────────────
def test_sanitize_features_forces_safe_color():
    bad = [{"name": "x", "type": "cds", "start": 0, "end": 5,
            "color": 'red" onmouseover="alert(1)'}]
    out = PL.sanitize_features(bad, 100)
    assert out[0]["color"] == PL._DEFAULT_COLOR        # injection scrubbed
    ok = PL.sanitize_features([{"start": 0, "end": 5, "color": "#6A89E4"}], 100)
    assert ok[0]["color"] == "#6A89E4"                  # valid hex kept


def test_sanitize_features_clamps_and_drops():
    feats = [
        {"start": -10, "end": 50},                      # start clamped to 0
        {"start": 90, "end": 9999},                     # end clamped to len
        {"start": 60, "end": 60},                       # empty → dropped
        {"start": "x", "end": 5},                       # non-int → dropped
        "not a dict",                                   # → dropped
    ]
    out = PL.sanitize_features(feats, 100)
    assert len(out) == 2
    assert out[0]["start"] == 0 and out[0]["end"] == 50
    assert out[1]["end"] == 100


def test_sanitize_features_caps_count(monkeypatch):
    monkeypatch.setattr(PL, "MAX_FEATURES", 10)
    out = PL.sanitize_features([{"start": 0, "end": 3} for _ in range(50)], 100)
    assert len(out) == 10


def test_sanitize_features_normalizes_strand():
    out = PL.sanitize_features([{"start": 0, "end": 5, "strand": "-"},
                                {"start": 0, "end": 5, "strand": 5}], 100)
    assert out[0]["strand"] == -1 and out[1]["strand"] == 1


# ── annotation ──────────────────────────────────────────────────────────────
def test_find_motifs_t7_promoter():
    feats = PL.find_motifs(SEQ)
    names = [f["name"] for f in feats]
    assert "T7 promoter" in names
    t7 = next(f for f in feats if f["name"] == "T7 promoter")
    assert t7["start"] == 0 and t7["strand"] == 1


def test_find_orfs():
    # ATG + 100 sense codons (no stop) + TAA
    orf = "ATG" + "GCT" * 100 + "TAA"
    feats = PL.find_orfs("CCCC" + orf + "CCCC", min_aa=50)
    assert any(f["type"] == "CDS" for f in feats)


def test_find_orfs_linear_on_stopless_input():
    # Audit DoS regression: a long, ATG-rich, stop-LESS sequence must return
    # fast (linear) and find nothing (no stop closes an ORF).
    import time
    seq = "ATG" * 40_000          # 120 kb, every codon an ATG, no stop
    t0 = time.time()
    feats = PL.find_orfs(seq, min_aa=50)
    assert feats == []
    assert time.time() - t0 < 2.0   # would be many seconds under the old O(n²)


def test_find_orfs_caps_count():
    # many short ORFs back-to-back → capped, doesn't explode
    unit = "ATG" + "GCT" * 60 + "TAA"      # ~61 aa ORF
    feats = PL.find_orfs(unit * 1200, min_aa=50, max_orfs=100)
    assert len(feats) <= 100


def test_auto_annotate_only_when_empty():
    p = PL.parse_plasmid(SEQ)
    PL.auto_annotate(p)
    assert p.get("auto_annotated") is True
    assert any(f["name"] == "T7 promoter" for f in p["features"])
    # second call is a no-op (features already present)
    before = list(p["features"])
    PL.auto_annotate(p)
    assert p["features"] == before


# ── GenBank round-trip ───────────────────────────────────────────────────
def test_genbank_roundtrip():
    p = PL.parse_plasmid(SEQ, topology="circular")
    p["name"] = "pTEST"
    p["features"] = [{"name": "T7 promoter", "type": "promoter", "start": 0,
                      "end": 18, "strand": 1, "color": "#5FA98A", "notes": ""}]
    gb = PL.build_genbank(p)
    assert "LOCUS" in gb and "promoter" in gb
    re_p = PL.parse_plasmid(gb)
    assert re_p["source"] == "genbank"
    assert re_p["sequence"] == SEQ
    assert re_p["topology"] == "circular"
    assert any(f["type"] == "promoter" for f in re_p["features"])


def test_build_fasta():
    p = PL.parse_plasmid(SEQ)
    fa = PL.build_fasta(p)
    assert fa.startswith(">")
    assert "".join(fa.split("\n")[1:]) == SEQ


# ── restriction + digest ─────────────────────────────────────────────────
def test_restriction_finds_ecori_sites():
    r = PL.restriction_analysis(SEQ, "circular", ["EcoRI", "BamHI"])
    ecori = next(e for e in r["enzymes"] if e["name"] == "EcoRI")
    assert ecori["count"] == 2
    assert "EcoRI" in r["double_cutters"]
    assert "BamHI" in r["non_cutters"]


def test_restriction_default_common_set():
    r = PL.restriction_analysis(SEQ, "circular")
    assert len(r["enzymes"]) == len(PL.COMMON_ENZYMES)


def test_fragments_linear_and_circular():
    # cuts at positions 3 and 23 in a 70-mer
    lin = PL._fragments([3, 23], 70, circular=False)
    assert sum(lin) == 70 and len(lin) == 3
    circ = PL._fragments([3, 23], 70, circular=True)
    assert sum(circ) == 70 and len(circ) == 2
    # single cut on a circle → one full-length linear fragment
    assert PL._fragments([10], 70, circular=True) == [70]
    # uncut → one fragment
    assert PL._fragments([], 70, circular=True) == [70]


def test_digest_two_cutter_circular_gives_two_fragments():
    d = PL.digest(SEQ, "circular", ["EcoRI"])
    assert d["n_fragments"] == 2
    assert sum(d["fragments"]) == len(SEQ)
    assert len(d["cut_positions"]) == 2


def test_digest_linear_one_cutter_gives_two_fragments():
    # SmaI? use EcoRI but linear: 2 sites → 3 fragments
    d = PL.digest(SEQ, "linear", ["EcoRI"])
    assert d["n_fragments"] == 3
    assert sum(d["fragments"]) == len(SEQ)


# ── Common-features auto-annotation library (2026-06) ────────────────────────
def test_common_features_detected_both_strands():
    """The expanded _MOTIFS library annotates common tags / sites / promoters
    on both strands, like SnapGene/Benchling."""
    parts = [
        "TAATACGACTCACTATAG",                  # T7 promoter (+)
        "ACGT" * 5,
        "ATAACTTCGTATAATGTATGCTATACGAAGTTAT",   # loxP
        "GGGG" * 4,
        "GATTACAAGGATGATGACGATAAG",             # FLAG
        "TT",
        "CATCACCATCACCATCAC",                   # 6xHis (CAT variant)
        "AAAACCCGGGTTT" * 3,
        "ACAAGTTTGTACAAAAAAGCAGGCT",            # attB1
    ]
    feats = PL.find_motifs("".join(parts))
    names = {f["name"] for f in feats}
    assert {"T7 promoter", "loxP", "FLAG tag", "6×His tag", "attB1"} <= names
    # reverse-strand detection
    rc = PL.revcomp("GAAGTTCCTATTCTCTAGAAAGTATAGGAACTTC")  # FRT
    rev = PL.find_motifs("CCCC" + rc + "GGGG")
    assert any(f["name"] == "FRT" and f.get("strand") == -1 for f in rev)