1 files changed, 64 insertions, 0 deletions
diff --git a/tests/test_phase3_flags.py b/tests/test_phase3_flags.py
new file mode 100644
index 0000000..4ba8be1
--- /dev/null
+++ b/tests/test_phase3_flags.py
@@ -0,0 +1,64 @@
+"""Tests for the Phase-3 flags: task_mode=is_prime, train_frac, lr_decay."""
+import torch
+
+from src.config import Config, parse_args
+from src.data import build_examples, decode_tokens, get_splits
+from src.eval import probe_report
+from src.model_api import build_model
+
+from tests.test_eval_classification import StubModel
+
+
+def _isprime_sieve(n):
+    """Primality via {2,3,5,7} with the d<c guard — EXACT on [2,100] (every composite <=100
+    has a prime factor <=10), wrong only on no-small-factor composites like 121."""
+    if n < 2:
+        return 0
+    return int(all(n % d != 0 for d in (2, 3, 5, 7) if d < n))
+
+
+def test_is_prime_examples():
+    cfg = Config(task_mode="is_prime")
+    ex = build_examples([2, 3, 4, 97, 100], cfg)
+    got = [decode_tokens(y, cfg) for _, y in ex]
+    assert got == [1, 1, 0, 1, 0], got
+
+
+def test_is_prime_sieve_stub_perfect_in_range():
+    """The {2,3,5,7} sieve is exact on [2,100] -> evaluate() must see EM 1.0."""
+    from src.train import evaluate
+    cfg = Config(task_mode="is_prime")
+    train_in, val_in = get_splits(cfg)
+    m = StubModel(_isprime_sieve)
+    for split in (train_in, val_in):
+        tok, em, _ = evaluate(m, build_examples(split, cfg), cfg)
+        assert em == 1.0 and tok == 1.0
+
+
+def test_is_prime_probe_classified_p1():
+    """On [101,200] the sieve errs exactly on {121,143,169,187} (classified prime) -> P1."""
+    cfg = Config(task_mode="is_prime")
+    r = probe_report(StubModel(_isprime_sieve), cfg)
+    assert sorted(e["n"] for e in r["errors"]) == [121, 143, 169, 187], r["errors"]
+    assert r["code"] == "P1", r["code"]
+
+
+def test_train_frac_subsamples_train_only():
+    cfg = Config(train_frac=0.5)
+    tr, va = get_splits(cfg)
+    assert len(va) == 30                      # val untouched (locked size)
+    assert len(tr) == 34                      # round(69*0.5) = 34 (banker's rounding)
+    assert not set(tr) & set(va)
+    # deterministic across calls
+    tr2, va2 = get_splits(Config(train_frac=0.5))
+    assert tr == tr2 and va == va2
+    # subsample is a subset of the full train split
+    tr_full, va_full = get_splits(Config(train_frac=1.0))
+    assert set(tr) <= set(tr_full)
+
+
+def test_new_flags_parse():
+    a = parse_args(["rnn", "0", "--task_mode", "is_prime", "--train_frac", "0.4", "--lr_decay", "True"])
+    assert a.task_mode == "is_prime" and a.train_frac == 0.4 and a.lr_decay is True
+    b = parse_args(["transformer", "0", "--lr_decay", "False"])
+    assert b.lr_decay is False and b.task_mode == "next_prime" and b.train_frac == 1.0