tests/test_preference_pairs.py

import json
import subprocess
import sys
from pathlib import Path


ROOT = Path(__file__).resolve().parent.parent
SCRIPT = ROOT / "scripts" / "build_preference_pairs.py"
OUT = ROOT / "data" / "test_pairs.jsonl"


REQUIRED_FIELDS = {
    "prompt", "chosen", "rejected", "scenario",
    "step", "issue_id", "is_attack", "kind",
}


def _run_builder(max_pairs: int = 20) -> None:
    if OUT.exists():
        OUT.unlink()
    cmd = [
        sys.executable, str(SCRIPT),
        "--max-pairs", str(max_pairs),
        "--out", str(OUT),
    ]
    res = subprocess.run(cmd, cwd=str(ROOT), capture_output=True, text=True)
    assert res.returncode == 0, f"builder failed: {res.stderr}\nSTDOUT:{res.stdout}"


def test_preference_pairs_smoke():
    _run_builder(max_pairs=20)
    assert OUT.exists(), f"output file not written: {OUT}"
    lines = OUT.read_text(encoding="utf-8").strip().splitlines()
    assert len(lines) == 20, f"expected 20 lines, got {len(lines)}"
    for i, line in enumerate(lines):
        rec = json.loads(line)
        missing = REQUIRED_FIELDS - rec.keys()
        assert not missing, f"line {i} missing fields: {missing}"
        assert rec["chosen"] != rec["rejected"], f"line {i} has chosen == rejected"
        assert isinstance(rec["prompt"], str) and rec["prompt"], f"line {i} empty prompt"
        assert rec["kind"] in {"contrastive_reasoning", "wrong_action_confident"}
        json.loads(rec["chosen"])
        json.loads(rec["rejected"])


if __name__ == "__main__":
    test_preference_pairs_smoke()
    print("PASS test_preference_pairs_smoke")