lean-bench/tests/test_pilot_scripts.py

"""Offline smoke test of the pilot scripts with a stub model (no network)."""
import importlib.util
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1] / "experiments" / "P0-pilot"


def load(name):
    spec = importlib.util.spec_from_file_location(name, ROOT / f"{name}.py")
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


class Stub:
    def complete_json(self, messages, **kw):
        if "roles" in messages[0]["content"] or "passage" in messages[0]["content"].lower() and "faithful" not in messages[0]["content"]:
            return {"roles": [
                {"kind": "object", "symbols": ["G"], "description": "finite group", "keywords": ["group", "finite"]},
                {"kind": "goal", "symbols": [], "description": "g^n = e", "keywords": ["identity"]},
            ]}
        return {"faithful": True, "confidence": 0.9, "reason": "ok"}


ITEM = {"id": "x1", "source_text": "Let G be a finite group. Then g^n = e.", "label": "t",
        "lean_code": "theorem t (G : Type*) [Group G] [Fintype G] (g : G) : g ^ Fintype.card G = 1 := sorry",
        "verified": True}


def test_run_item_and_score(tmp_path):
    rp, sc = load("run_pilot"), load("score")
    row = rp.run_item(Stub(), ITEM)
    assert row["id"] == "x1" and row["compiled"] and "judge" in row
    assert set(sc.flags(row)) == {"compile_only", "lint", "role", "judge"}
    rows = [dict(row, verdict=v) for v in ("wrong", "faithful", "faithful", "wrong")]
    out = sc.report(rows, lambda r: r["verdict"] == "wrong")
    assert out["judge"]["n"] == 4
    assert sc.decide(out) in {"GO", "PARTIAL", "NO-GO"}


def test_complementary_signal_is_partial():
    sc = load("score")
    rows = [
        {"compiled": True, "lint": [], "role": {"flagged": i < 3, "score": 0.1},
         "judge": {"faithful": False if 3 <= i < 8 else True}, "verdict": "wrong"}
        for i in range(10)
    ] + [
        {"compiled": True, "lint": [], "role": {"flagged": False, "score": 1},
         "judge": {"faithful": True}, "verdict": "faithful"}
    ]
    summary = sc.report(rows, lambda row: row["verdict"] == "wrong")
    assert summary["role"]["recall"] == 0.3
    assert summary["role"]["complementary_wrong"] == 3
    assert sc.decide(summary) == "PARTIAL"


def test_score_rejects_missing_labels_and_detector_failures(tmp_path):
    import argparse
    import json

    import pytest

    sc = load("score")
    args = argparse.Namespace(results=tmp_path / "results.jsonl", labels=tmp_path / "labels.json",
                              key=tmp_path / "key.json")
    args.key.write_text(json.dumps({"key": {"P01": "x1"}}))
    args.labels.write_text(json.dumps({"labels": {"P01": {"verdict": None}}}))
    args.results.write_text(json.dumps({"id": "x1"}) + "\n")
    with pytest.raises(ValueError, match="verdict"):
        sc.load(args)
    args.labels.write_text(json.dumps({"labels": {"P01": {"verdict": "wrong"}}}))
    with pytest.raises(ValueError, match="detector failure"):
        sc.load(args)
    args.results.write_text("")
    with pytest.raises(ValueError, match="every keyed item"):
        sc.load(args)


def test_request_budget_survives_resume_and_transport_retries(tmp_path, monkeypatch):
    import argparse
    import json
    import urllib.error

    import pytest

    from leanbench.config import Settings

    rp = load("run_pilot")
    args = argparse.Namespace(items=tmp_path / "items.jsonl", out=tmp_path / "results.jsonl",
                              max_calls=2, dry_run=False)
    args.items.write_text(json.dumps(ITEM) + "\n")
    monkeypatch.setenv("DASHSCOPE_API_KEY", "test-only")
    settings = Settings()
    settings.llm.max_retries = 5
    requests = []

    def fail(*args, **kwargs):
        requests.append(1)
        raise urllib.error.URLError("offline test")

    monkeypatch.setattr("urllib.request.urlopen", fail)
    monkeypatch.setattr("leanbench.llm.time.sleep", lambda *_: None)
    assert rp.execute(args, [ITEM], settings) == 2
    assert len(requests) == 2
    assert rp.execute(args, [ITEM], settings) == 2
    assert len(requests) == 2
    meta = json.loads(args.out.with_suffix(".meta.json").read_text())
    assert meta["request_attempts"] == 2
    assert len(meta["prompt_hashes"]["judge"]) == 64
    assert meta["status"] == "budget_exhausted"
    settings.llm.model = "other-model"
    with pytest.raises(ValueError, match="identity changed"):
        rp.execute(args, [ITEM], settings)


def test_run_rejects_missing_source_passage():
    import pytest

    rp = load("run_pilot")
    with pytest.raises(ValueError, match="Missing source_text"):
        rp.run_item(Stub(), {**ITEM, "source_text": ""})


def test_restore_passages_checks_cohort_and_normalizes_archived_imports():
    import pytest

    prepare = load("prepare_items")
    item = {"id": "x", "lean_code": "example : True := trivial", "imports": ["Mathlib"]}
    original = {**item, "imports": '["Mathlib"]', "source_text": "A private passage."}
    restored = prepare.restore([item], {"items": [original]})
    assert restored == [{**item, "source_text": "A private passage."}]
    with pytest.raises(ValueError, match="does not match"):
        prepare.restore([{**item, "lean_code": "changed"}], {"items": [original]})


def test_judge_rejects_string_boolean():
    import pytest

    from leanbench.llm import ScriptedClient

    rp = load("run_pilot")
    with pytest.raises(ValueError, match="unusable judge"):
        rp.judge(ScriptedClient(['{"faithful": "false"}']), "passage", "code")