Blame for tests/test_learning.py

"""Apprentissage automatique : journalisation, format d'invite partagé, job LoRA (sur un Qwen2 minuscule, CPU)."""
import importlib.util
import json
from pathlib import Path

import pytest

from mouh.jev.types import JevMessageType as T, TaskPayload
from mouh.learning import TrajectoryRecorder, build_record, chat_messages, render_files, split_for, task_hash
from mouh.learning import prompts as P
from mouh.verification import CallableGenerator, TestTimeVerificationLoop
from mouh.verification.generators import parse_files

ROOT = Path(__file__).resolve().parents[1]
TESTS = {"test_solution.py": "from solution import add\n\ndef test_add():\n    assert add(2, 3) == 5\n"}


def load_job():
    spec = importlib.util.spec_from_file_location("train_unique", ROOT / "jobs" / "train_unique.py")
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


JOB = load_job()


# ---------------------------------------------------------------- format partagé Space <-> job
@pytest.mark.parametrize("prev,diags", [(None, None), ({"solution.py": "x=1\n"}, ["boom", "  autre  "])])
def test_job_prompt_copy_is_identical(prev, diags):
    assert JOB.SYSTEM_PROMPT == P.SYSTEM_PROMPT
    assert JOB.chat_messages("t", TESTS, prev, diags) == P.chat_messages("t", TESTS, prev, diags)
    files = {"solution.py": "def add(a, b):\n    return a + b\n", "pkg/util.py": "X = 1\n"}
    assert JOB.render_files(files) == P.render_files(files)


def test_rendered_target_parses_back():
    files = {"solution.py": "def add(a, b):\n    return a + b\n", "pkg/util.py": "X = 1\n"}
    assert parse_files(render_files(files)) == files
    assert JOB.parse_files(render_files(files)) == files


# ---------------------------------------------------------------- bus et enregistrement
def run_with_repair():
    calls = []

    def gen(req):
        calls.append(req)
        return {"solution.py": "def add(a, b):\n    return a + b\n"}

    loop = TestTimeVerificationLoop(CallableGenerator(gen), max_attempts=3)
    task = TaskPayload(instruction="Écrire add", candidate_files={"solution.py": "def add(a, b):\n    return a - b\n"},
                       test_files=TESTS)
    return task, loop.run(task), calls


def test_router_decision_precedes_code_proposal():
    _, out, calls = run_with_repair()
    types = [m.type for m in out.messages]
    for i, t in enumerate(types):
        if t == T.CODE_PROPOSAL:
            assert types[i - 1] == T.ROUTE_DECISION, types
    assert calls and calls[0].experts  # le générateur reçoit les experts décidés par le routeur


def test_build_record_keeps_attempts_and_repair():
    task, out, _ = run_with_repair()
    rec = build_record(task, out, "test")
    assert rec["certified"] and rec["status"] == "certified"
    assert [a["certified"] for a in rec["attempts"]] == [False, True]
    assert rec["attempts"][0]["diagnostics"] and "a - b" in rec["attempts"][0]["files"]["solution.py"]
    assert rec["final_files"]["solution.py"].endswith("a + b\n")
    assert rec["split"] == split_for(rec["task_hash"]) and rec["task_hash"] == task_hash(task.instruction, TESTS)
    json.dumps(rec)  # sérialisable


def test_task_hash_ignores_whitespace_and_split_is_deterministic():
    assert task_hash("a  b", {"t.py": "x =  1"}) == task_hash("a b", {"t.py": "x = 1\n"})
    splits = {split_for(task_hash(f"t{i}", TESTS)) for i in range(200)}
    assert splits == {"train", "heldout"}


class FakeScheduler:
    def __init__(self, **kw):
        import threading

        self.kw, self.lock = kw, threading.Lock()


def test_recorder_local_and_remote(tmp_path):
    local = TrajectoryRecorder(tmp_path / "a")
    assert not local.remote
    local.record({"x": 1})
    assert json.loads(local.path.read_text()) == {"x": 1}
    remote = TrajectoryRecorder(tmp_path / "b", "me/ds", "tok", every_min=3, scheduler_factory=FakeScheduler)
    assert remote.remote and remote._scheduler.kw["private"] and remote._scheduler.kw["repo_type"] == "dataset"
    assert remote._scheduler.kw["allow_patterns"] == ["data/*.jsonl"]


def test_service_records_session(tmp_path, monkeypatch):
    import mouh.learning.recorder as R
    from mouh.service import stream_session

    monkeypatch.setattr(R, "_RECORDER", TrajectoryRecorder(tmp_path))
    snaps = list(stream_session("add", "def add(a, b):\n    return a + b\n", TESTS["test_solution.py"]))
    rep = json.loads(snaps[-1]["report"])
    assert rep["journalisation"].startswith("oui")
    rec = json.loads(R._RECORDER.path.read_text().splitlines()[-1])
    assert rec["certified"] and rec["schema"] == "mouh-session/1"
    snaps = list(stream_session("add", "def add(a, b):\n    return a + b\n", TESTS["test_solution.py"], record=False))
    assert json.loads(snaps[-1]["report"])["journalisation"].startswith("non")
    assert len(R._RECORDER.path.read_text().splitlines()) == 1


# ---------------------------------------------------------------- données du job
def session(i, *, certified=True, tests=None, split=None, failed_first=True):
    tests = tests or {"test_solution.py": f"from solution import f\n\ndef test_f():\n    assert f() == {i}\n"}
    th = task_hash(f"tâche {i}", tests)
    final = {"solution.py": f"def f():\n    return {i}\n"}
    attempts = []
    if failed_first:
        attempts.append({"attempt": 1, "origin": "user", "files": {"solution.py": "def f():\n    return -1\n"},
                         "score": 0.7, "certified": False, "diagnostics": [f"assert -1 == {i}"]})
    attempts.append({"attempt": len(attempts) + 1, "origin": "model", "files": final, "score": 1.0,
                     "certified": certified, "diagnostics": []})
    return {"schema": "mouh-session/1", "session_id": f"s{i}", "task_hash": th, "split": split or split_for(th),
            "instruction": f"tâche {i}", "tests": tests, "proofs": [], "generator": "x",
            "status": "certified" if certified else "exhausted", "certified": certified,
            "final_files": final if certified else None, "attempts": attempts}


def test_build_examples_filters_and_separates_heldout():
    sessions = [session(1, split="train"), session(2, split="heldout"), session(3, certified=False, split="train"),
                session(4, split="train", tests={"test_solution.py": "def test_x():\n    assert True\n"}),
                session(1, split="train")]  # doublon exact
    train, heldout, stats = JOB.build_examples(sessions)
    assert stats["kept"] == 3 and [t["task_hash"] for t in heldout] == [sessions[1]["task_hash"]]
    assert sorted(e["kind"] for e in train) == ["fix", "gen"]  # doublon retiré, tâche 4 rejetée (n'importe pas solution)
    assert all(sessions[1]["instruction"] not in json.dumps(e) for e in train)


def test_run_tests_really_executes():
    t = {"test_solution.py": "from solution import f\n\ndef test_f():\n    assert f() == 3\n"}
    assert JOB.run_tests({"solution.py": "def f():\n    return 3\n"}, t)
    assert not JOB.run_tests({"solution.py": "def f():\n    return 4\n"}, t)
    assert not JOB.run_tests({"solution.py": "import time\ndef f():\n    time.sleep(5)\n"}, t, timeout_s=2)


# ---------------------------------------------------------------- job complet sur un modèle minuscule
@pytest.fixture(scope="module")
def tiny_qwen(tmp_path_factory):
    pytest.importorskip("peft")
    from tokenizers import Tokenizer, decoders, models, pre_tokenizers, trainers
    from transformers import PreTrainedTokenizerFast, Qwen2Config, Qwen2ForCausalLM

    d = tmp_path_factory.mktemp("tiny_qwen")
    tk = Tokenizer(models.BPE())
    tk.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tk.decoder = decoders.ByteLevel()
    corpus = [P.SYSTEM_PROMPT, render_files({"solution.py": "def f():\n    return 1\n"}), "## Tâche ## Tests assert"] * 20
    tk.train_from_iterator(corpus, trainers.BpeTrainer(vocab_size=400, special_tokens=["<|endoftext|>", "<|im_start|>", "<|im_end|>"],
                                                        initial_alphabet=pre_tokenizers.ByteLevel.alphabet()))
    tok = PreTrainedTokenizerFast(tokenizer_object=tk, eos_token="<|im_end|>", pad_token="<|endoftext|>")
    tok.chat_template = ("{% for m in messages %}<|im_start|>{{ m['role'] }}\n{{ m['content'] }}<|im_end|>\n{% endfor %}"
                         "{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}")
    tok.save_pretrained(d)
    cfg = Qwen2Config(vocab_size=len(tok), hidden_size=32, intermediate_size=64, num_hidden_layers=2,
                      num_attention_heads=4, num_key_value_heads=2, max_position_embeddings=4096,
                      eos_token_id=tok.eos_token_id, pad_token_id=tok.pad_token_id, tie_word_embeddings=True)
    Qwen2ForCausalLM(cfg).save_pretrained(d)
    return str(d)


def write_sessions(path: Path, n: int, start: int = 0):
    path.mkdir(parents=True, exist_ok=True)
    with open(path / f"s{start}.jsonl", "w") as fh:
        for i in range(start, start + n):
            fh.write(json.dumps(session(i)) + "\n")


def test_job_full_cycle_offline(tiny_qwen, tmp_path):
    data, out = tmp_path / "data", tmp_path / "out"
    write_sessions(data, 40)
    args = ["--local-data", str(data), "--local-out", str(out), "--base-model", tiny_qwen, "--min-new-examples", "10",
            "--min-heldout", "1", "--epochs", "1", "--max-new-tokens", "8", "--grad-accum", "4", "--max-len", "1024"]
    assert JOB.main(args) == 0
    state = json.loads((out / "state.json").read_text())
    run = json.loads(next((out / "runs").glob("*.json")).read_text())
    assert state["trained_on"] == run["data"]["train_examples"] > 0
    assert run["train"]["optimizer_steps"] > 0 and "base" in run and "candidate" in run
    assert run["promoted"] is False  # un modèle aléatoire ne fait passer aucun test : rien n'est promu
    assert (out / "adapter_candidate" / "adapter_config.json").exists() and not (out / "adapter").exists()
    # second passage sans nouvelles données : le job s'arrête avant de charger quoi que ce soit
    assert JOB.main(args) == 0
    assert len(list((out / "runs").glob("*.json"))) == 1


def test_job_promotes_only_when_strictly_better(tiny_qwen, tmp_path, monkeypatch):
    data, out = tmp_path / "data", tmp_path / "out"
    write_sessions(data, 40)
    scores = iter([1, 3, 3, 3, 2])  # cycle 1 : base 1, candidat 3 -> promu ; cycle 2 : base 3, actuel 3, candidat 2 -> non
    monkeypatch.setattr(JOB, "pass_at_1", lambda *a, **k: {"passed": next(scores), "n": 5, "rate": 0, "ci95": [0, 1]})
    args = ["--local-data", str(data), "--local-out", str(out), "--base-model", tiny_qwen, "--min-new-examples", "1",
            "--min-heldout", "1", "--epochs", "1", "--max-len", "1024"]
    assert JOB.main(args) == 0
    assert (out / "adapter" / "adapter_config.json").exists()
    assert json.loads((out / "state.json").read_text())["promoted"]["candidate"]["passed"] == 3
    write_sessions(data, 10, start=100)
    assert JOB.main(args) == 0
    hist = json.loads((out / "state.json").read_text())["history"]
    assert [h["promoted"] for h in hist] == [True, False]


def test_unique_generator_runs_in_the_loop(tiny_qwen, monkeypatch):
    import mouh.service as S

    monkeypatch.setenv("UNIQUE_BASE_MODEL", tiny_qwen)
    monkeypatch.setenv("UNIQUE_MAX_NEW_TOKENS", "8")
    monkeypatch.delenv("UNIQUE_ADAPTER_REPO", raising=False)
    S._MODEL_CACHE.pop("unique", None)
    snaps = list(S.stream_session("add", "def add(a, b):\n    return a - b\n", TESTS["test_solution.py"], max_attempts=2,
                                  generator_choice=S.GENERATOR_CHOICES[2], record=False))
    assert snaps[-1]["status"].startswith("AUCUNE SORTIE")  # modèle aléatoire : rien de non prouvé n'est émis
    S._MODEL_CACHE.pop("unique", None)


# ---------------------------------------------------------------- pilote automatique (tout sur HF)
class FakeJobsApi:
    def __init__(self):
        import types

        self.jobs, self.created, self.triggered = [], [], []
        self._ns = types.SimpleNamespace

    def list_scheduled_jobs(self, labels=None):
        return [j for j in self.jobs if all(j.labels.get(k) == v for k, v in (labels or {}).items())]

    def create_scheduled_uv_job(self, script, **kw):
        assert Path(script).exists() and kw["secrets"]["HF_TOKEN"] == "tok"
        job = self._ns(id=f"j{len(self.jobs)}", schedule=kw["schedule"], labels=kw["labels"],
                       status=self._ns(next_job_run_at="lundi", last_job=None))
        self.jobs.append(job)
        self.created.append(kw)
        return job

    def trigger_scheduled_job(self, scheduled_job_id):
        self.triggered.append(scheduled_job_id)
        return self._ns(url=f"https://hf.co/jobs/{scheduled_job_id}")

    def hf_hub_download(self, repo, name):
        raise FileNotFoundError


def test_autopilot_is_opt_in_and_idempotent(monkeypatch):
    from mouh.learning import autopilot as A

    api = FakeJobsApi()
    monkeypatch.setenv("HF_TOKEN", "tok")
    monkeypatch.delenv("UNIQUE_AUTOTRAIN", raising=False)
    msg, job = A.ensure_schedule(api, A.settings())
    assert job is None and "désactivé" in msg and not api.created  # aucune dépense sans accord explicite
    monkeypatch.setenv("UNIQUE_AUTOTRAIN", "1")
    _, j1 = A.ensure_schedule(api, A.settings())
    _, j2 = A.ensure_schedule(api, A.settings())
    assert j1.id == j2.id and len(api.created) == 1
    assert api.created[0]["flavor"] == "t4-small" and api.created[0]["env"]["MOUH_DATASET_REPO"]
    assert "j0" in A.status_markdown(api) and "Aucun cycle" in A.status_markdown(api)


def test_trigger_requires_admin_key(monkeypatch):
    from mouh.learning import autopilot as A

    api = FakeJobsApi()
    monkeypatch.setenv("HF_TOKEN", "tok")
    monkeypatch.setenv("UNIQUE_AUTOTRAIN", "1")
    monkeypatch.delenv("UNIQUE_ADMIN_KEY", raising=False)
    assert "désactivé" in A.trigger_now("x", api) and not api.triggered
    monkeypatch.setenv("UNIQUE_ADMIN_KEY", "secret")
    assert "incorrecte" in A.trigger_now("mauvaise", api) and not api.triggered
    assert "Cycle lancé" in A.trigger_now("secret", api) and api.triggered == ["j0"]


def test_app_builds():
    import importlib

    import app  # noqa: F401 — l'interface se construit sans réseau ni jeton

    importlib.reload(app)