tests/test_clone_sim.py

"""Unit tests for the cloning simulator (dee/core/cloning.py)."""
import pytest

from dee.core import cloning as C


# ── Gibson / overlap assembly ───────────────────────────────────────────────
def test_gibson_native_overlaps_circular():
    # three fragments sharing 20-bp overlaps, arranged head-to-tail in a circle.
    o1 = "ACGTACGTACGTACGTACGT"          # A↔B overlap
    o2 = "TTTTGGGGCCCCAAAATTTT"          # B↔C overlap
    o3 = "GGGGCCCCAAAATTTTGGGG"          # C↔A overlap
    A = o3 + "AAAAAAAAAAAA" + o1
    B = o1 + "CCCCCCCCCCCC" + o2
    Cc = o2 + "GGGGGGGGGGGG" + o3
    out = C.gibson([{"name": "A", "seq": A}, {"name": "B", "seq": B}, {"name": "C", "seq": Cc}])
    assert out["ok"] is True
    assert out["topology"] == "circular"
    assert all(j["kind"] == "native" for j in out["junctions"])
    assert out["length"] == len(A) + len(B) + len(Cc) - 60   # minus three 20-bp overlaps


def test_gibson_auto_orders_fragments():
    o1 = "ACGTACGTACGTACGTACGT"
    o2 = "TTTTGGGGCCCCAAAATTTT"
    A = "AAAAAAAA" + o1
    B = o1 + "CCCCCCCC" + o2
    Cc = o2 + "GGGGGGGG"
    out = C.gibson([{"name": "C", "seq": Cc}, {"name": "A", "seq": A}, {"name": "B", "seq": B}],
                   circular=False)
    assert out["ok"] is True
    assert out["order"] == ["A", "B", "C"]
    assert out["length"] == len(A) + len(B) + len(Cc) - 40


def test_gibson_designs_overlaps_and_primers_for_bare_fragments():
    A = "ATGGTGAGCAAGGGCGAGGAGCTGTTCACC"
    B = "GACGAGCTGTACAAGTAACTCGAGCTAGCAT"
    out = C.gibson([{"name": "A", "seq": A}, {"name": "B", "seq": B}], circular=False)
    assert out["ok"] is True
    assert any(j["kind"] == "designed" for j in out["junctions"])
    assert out["length"] == len(A) + len(B)
    prim = {p["fragment"]: p for p in out["primers"]}
    assert prim["B"]["forward_tail"] > 0
    assert prim["B"]["forward"].startswith(A[-25:])


def test_gibson_needs_two_fragments():
    assert C.gibson([{"name": "x", "seq": "ACGT"}])["ok"] is False


# ── Golden Gate ──────────────────────────────────────────────────────────────
def _bsa_part(left_oh, body, right_oh):
    return "GAATT" + "GGTCTC" + "T" + left_oh + body + right_oh + "T" + "GAGACC" + "TTAAG"


def test_excise_part_reads_overhangs():
    p = C._excise_part(_bsa_part("AAAA", "GCTAGCGGGG", "CCCC"), "BsaI")
    assert p["left_oh"] == "AAAA"
    assert p["right_oh"] == "CCCC"
    assert p["inner"] == "AAAA" + "GCTAGCGGGG" + "CCCC"


def test_golden_gate_assembles_circle():
    p1 = _bsa_part("AAAA", "GCTAGCAAAA", "CCCC")
    p2 = _bsa_part("CCCC", "TTGGCCAATT", "AAAA")
    out = C.golden_gate([{"name": "P1", "seq": p1}, {"name": "P2", "seq": p2}], enzyme="BsaI")
    assert out["ok"] is True
    assert out["topology"] == "circular"
    assert out["length"] == (4 + len("GCTAGCAAAA")) + (4 + len("TTGGCCAATT"))
    assert out["order"] == ["P1", "P2"]


def test_golden_gate_rejects_mismatched_overhangs():
    p1 = _bsa_part("AAAA", "GCTAGCAAAA", "CCCC")
    p2 = _bsa_part("GGGG", "TTGGCCAATT", "AAAA")
    out = C.golden_gate([{"name": "P1", "seq": p1}, {"name": "P2", "seq": p2}], enzyme="BsaI")
    assert out["ok"] is False


def test_golden_gate_rejects_non_cassette():
    out = C.golden_gate([{"name": "P1", "seq": "ACGTACGTACGT"}, {"name": "P2", "seq": "TTTTGGGG"}])
    assert out["ok"] is False


# ── restriction cut-and-ligate ──────────────────────────────────────────────
def test_restriction_clone_single_enzyme():
    vector = "AAAACCCC" + "GAATTC" + "TTTTGGGGCCCCAAAA"
    insert = "CC" + "GAATTC" + "ATGCATGCATGCATGC" + "GAATTC" + "GG"
    out = C.restriction_clone(vector, insert, "EcoRI", vector_topology="circular")
    assert out["ok"] is True
    assert out["topology"] == "circular"
    assert "ATGCATGCATGCATGC" in out["assembled"]
    assert out["length"] > len(vector)


def test_restriction_clone_linear_splices_the_insert_in_at_the_cut():
    # A single cut on a LINEAR vector produces two loose ends, not a seam to
    # close — the correct operation is splicing the insert IN at the cut
    # (left arm, insert, right arm), not gluing the two arms back together
    # unchanged and appending the insert onto the end (the original bug: a
    # silent "ok": True carrying the untouched vector plus a dangling tail,
    # mislabelled "topology": "circular" on top of that).
    left, right = "AAAACCCC", "TTTTGGGGCCCCAAAA"
    vector = left + "GAATTC" + right
    insert = "CC" + "GAATTC" + "ATGCATGCATGCATGC" + "GAATTC" + "GG"
    out = C.restriction_clone(vector, insert, "EcoRI", vector_topology="linear")
    assert out["ok"] is True
    assert out["topology"] == "linear"
    # The insert lands BETWEEN the two original arms — not after them.
    assert out["assembled"].startswith(left)
    assert out["assembled"].endswith(right)
    assert "ATGCATGCATGCATGC" in out["assembled"]
    insert_start = out["assembled"].index("ATGCATGCATGCATGC")
    assert insert_start < len(out["assembled"]) - len(right)
    # Nothing lost, nothing silently kept whole: length is vector + insert
    # fragment exactly, not vector + insert (no-op backbone) or vector alone.
    i0 = insert.index("GAATTC")
    i1 = insert.rindex("GAATTC")
    insert_frag_len = i1 - i0
    assert out["length"] == len(vector) + insert_frag_len
    assert len(out["features"]) == 3


def test_restriction_clone_requires_unique_vector_site():
    vector = "GAATTCAAAAGAATTC"
    insert = "GAATTCATGCGAATTC"
    out = C.restriction_clone(vector, insert, "EcoRI", vector_topology="linear")
    assert out["ok"] is False


def test_restriction_clone_requires_flanked_insert():
    vector = "AAAAGAATTCTTTT"
    insert = "ATGCATGCATGC"
    out = C.restriction_clone(vector, insert, "EcoRI")
    assert out["ok"] is False


# ── Yeast in-vivo (homologous-recombination) assembly ───────────────────────
import random as _random


def _ydna(rng, n):
    return "".join(rng.choice("ACGT") for _ in range(n))


def test_yeast_designs_arms_primers_and_plan():
    rng = _random.Random(11)
    vec, gA, gB = _ydna(rng, 240), _ydna(rng, 180), _ydna(rng, 180)
    frags = [{"name": "pRS416 URA3 vector", "seq": vec},
             {"name": "geneA", "seq": gA}, {"name": "geneB", "seq": gB}]
    r = C.yeast_assembly(frags, circular=True)
    assert r["ok"] and r["method"] == "yeast_hr" and r["topology"] == "circular"
    assert r["order"] == ["pRS416 URA3 vector", "geneA", "geneB"]
    # distinct fragments → every junction is a designed, UNIQUE homology arm
    assert all(j["kind"] == "designed" for j in r["junctions"])
    assert all(j["unique"] for j in r["junctions"])
    assert not any("share homology" in w for w in r["warnings"])
    # chimeric primers carry a 40-bp homology tail
    assert any(p["forward_tail"] == C.YEAST_ARM for p in r["primers"])
    # marker detected from the vector name; plan + one verification pair per junction
    assert r["marker"] == "pRS416 URA3 vector"
    assert set(r["plan"]) == {"transform", "selection", "verify"}
    assert len(r["verification"]) == len(r["junctions"])
    assert all(v["forward"] and v["reverse"] for v in r["verification"])


def test_yeast_flags_repeated_terminator_as_misassembly_risk():
    rng = _random.Random(12)
    term = _ydna(rng, 45)                       # a terminator reused on two genes
    vec = _ydna(rng, 220)
    gA = _ydna(rng, 150) + term
    gB = _ydna(rng, 150) + term                 # SAME terminator at the 3' end
    frags = [{"name": "vec LEU2", "seq": vec},
             {"name": "geneA-tCYC1", "seq": gA}, {"name": "geneB-tCYC1", "seq": gB}]
    r = C.yeast_assembly(frags, circular=True)
    assert r["ok"]
    # the two junctions reading off the repeated terminator share homology
    assert sum(1 for j in r["junctions"] if not j["unique"]) >= 2
    assert any("share homology" in w for w in r["warnings"])


def test_yeast_missing_marker_warns_and_can_be_named():
    rng = _random.Random(13)
    frags = [{"name": "backbone", "seq": _ydna(rng, 200)},
             {"name": "insert", "seq": _ydna(rng, 200)}]
    r = C.yeast_assembly(frags, circular=True)
    assert r["marker"] is None
    assert any("selectable marker" in w for w in r["warnings"])
    r2 = C.yeast_assembly(frags, circular=True, marker="HIS3")
    assert r2["marker"] == "HIS3"
    assert not any("selectable marker" in w for w in r2["warnings"])


# ── Yeast chromosomal integration (knock-in / knock-out) ────────────────────
def test_yeast_integration_designs_cassette_arms_and_verification():
    rng = _random.Random(21)
    up = _ydna(rng, 300)          # genomic sequence 5' of the site
    dn = _ydna(rng, 300)          # genomic sequence 3' of the site
    cassette = _ydna(rng, 1400)   # a marker cassette to drop in
    r = C.yeast_integration(up, dn, cassette, marker="KanMX")
    assert r["ok"] and r["method"] == "yeast_integration" and r["topology"] == "linear"
    # the integrated allele is flank + cassette + flank, nothing lost
    assert r["assembled"] == up + cassette + dn
    assert r["length"] == len(up) + len(cassette) + len(dn)
    # the cassette is PCR'd with genomic-homology tails: fwd carries up's 3' arm
    p = r["cassette_primers"][0]
    assert p["forward_tail"] == r["arm"] == C.YEAST_ARM
    assert p["forward"].startswith(up[-C.YEAST_ARM:])
    assert p["reverse"].startswith(C.revcomp(dn[:C.YEAST_ARM]))
    # two designed junctions, both unique for distinct random flanks
    assert len(r["junctions"]) == 2
    assert all(j["kind"] == "designed" and j["unique"] for j in r["junctions"])
    assert r["marker"] == "KanMX"
    assert set(r["plan"]) == {"transform", "selection", "verify"}


def test_yeast_integration_verification_primer_sits_outside_the_arm():
    """On-target confirmation: the 5' check primer anneals in the flanking
    genome OUTSIDE the homology arm, and the 3' check primer OUTSIDE the 3'
    arm. A primer inside the arm (or inside the cassette) would give a band at
    any integration site and prove nothing."""
    rng = _random.Random(22)
    up, dn, cassette = _ydna(rng, 300), _ydna(rng, 300), _ydna(rng, 1200)
    r = C.yeast_integration(up, dn, cassette, marker="HIS3")
    v5, v3 = r["verification"]
    arm = r["arm"]
    # the 5' forward primer is sourced from genome 5' of the arm
    assert v5["forward"] and v5["forward"] in up[:-arm]
    assert v5["forward"] not in up[-arm:]
    # the 3' reverse primer anneals to genome 3' of the arm (bottom strand)
    assert v3["reverse"] and C.revcomp(v3["reverse"]) in dn[arm:]
    assert "OUTSIDE" in v5["note"] and "OUTSIDE" in v3["note"]


def test_yeast_integration_flags_arms_that_share_homology():
    rng = _random.Random(23)
    shared = _ydna(rng, 26)                 # >= 20 bp present in BOTH arms
    up = _ydna(rng, 220) + shared           # up[-40:] ends with `shared`
    dn = shared + _ydna(rng, 220)           # dn[:40]  starts with `shared`
    r = C.yeast_integration(up, dn, _ydna(rng, 800), marker="URA3")
    assert r["ok"]
    assert all(j["unique"] is False for j in r["junctions"])
    assert any("share sequence" in w for w in r["warnings"])


def test_yeast_integration_always_says_to_blast_the_arms():
    """The engine cannot see the genome, so off-target risk is never silently
    cleared — it is always surfaced as a thing the user must check."""
    rng = _random.Random(24)
    r = C.yeast_integration(_ydna(rng, 250), _ydna(rng, 250), _ydna(rng, 900), marker="LEU2")
    assert any("BLAST" in w for w in r["warnings"])


def test_yeast_integration_rejects_flanks_shorter_than_the_arm():
    rng = _random.Random(25)
    r = C.yeast_integration(_ydna(rng, 20), _ydna(rng, 250), _ydna(rng, 400))
    assert r["ok"] is False
    assert "homology arm" in r["error"]


def test_yeast_integration_warns_when_no_marker_named():
    rng = _random.Random(26)
    r = C.yeast_integration(_ydna(rng, 250), _ydna(rng, 250), _ydna(rng, 400))
    assert r["marker"] is None
    assert any("marker" in w for w in r["warnings"])