code/training/tests/test_data_pipeline.py

#!/usr/bin/env python3
"""Offline end-to-end test for the data-engineering layer.

Runs with stdlib + PyYAML only (no datasets/torch), so it is safe on any host:

    python training/tests/test_data_pipeline.py

It exercises:
  1. each adapter on rows that mirror the *real* HF schemas (verified 2026-06-27)
  2. build_sft_dataset.py over fixture raw splits  (adapt -> route -> dedup -> card)
  3. synthesize_think.py --mock over the to_synthesize queue
  4. build_sft_dataset.py --include-synthesized folding the traces back in
  5. split_jsonl.py producing train/val
  6. the run_sft assistant-<think> invariant on every ready row
"""

from __future__ import annotations

import json
import subprocess
import sys
import tempfile
from pathlib import Path

SCRIPTS = Path(__file__).resolve().parents[1] / "scripts"
CONFIGS = Path(__file__).resolve().parents[1] / "configs"
sys.path.insert(0, str(SCRIPTS))

import sft_adapters as A  # noqa: E402

RESULTS: list[str] = []


def check(name: str, condition: bool, detail: str = "") -> None:
    if condition:
        RESULTS.append(f"PASS {name}")
    else:
        RESULTS.append(f"FAIL {name} :: {detail}")
        raise AssertionError(f"{name} :: {detail}")


def has_think(content: str) -> bool:
    return "<think>" in content and "</think>" in content


def assistant_contents(messages):
    return [m["content"] for m in messages if m["role"] == "assistant"]


# --------------------------------------------------------------------------- #
# fixtures mirroring real HF schemas
# --------------------------------------------------------------------------- #
FIXTURES: dict[str, list[dict]] = {
    "primevul": [
        {"idx": 0, "project": "openssl", "target": 1, "func": "long ssl_get_algorithm2(SSL *s){ return s->x; }",
         "cwe": "[\"CWE-310\"]", "cve": "CVE-2013-6449", "cve_desc": "Obtains a version number from an incorrect source."},
        {"idx": 1, "project": "curl", "target": 0, "func": "int add(int a, int b){ return a + b; }",
         "cwe": "[]", "cve": "", "cve_desc": ""},
    ],
    "diversevul": [
        {"func": "static boolean ReadICCProfile(j_decompress_ptr p){ char magick[12]; return 1; }",
         "target": 1, "cwe": ["CWE-416"], "project": "ImageMagick", "message": "uaf"},
    ],
    "megavul": [
        {"instruction": "Analyze the following code function for security vulnerabilities",
         "input": "void f(char *s){ char b[8]; strcpy(b, s); }",
         "output": "Vulnerability Detected:\nStack buffer overflow via strcpy into a fixed buffer.",
         "cwe_ids": ["CWE-787"], "cve_id": "CVE-0000-1", "is_vulnerable": 1},
        {"instruction": "Analyze the following code function for security vulnerabilities",
         "input": "int g(void){ return 0; }",
         "output": "No Vulnerability Detected:\nThis code is secure.",
         "cwe_ids": [], "cve_id": "", "is_vulnerable": 0},
    ],
    "crossvul": [
        {"cwe_id": "CWE-89", "cwe_description": "SQL Injection", "language": "php",
         "vulnerable_code": "$q = \"SELECT * FROM u WHERE id=\".$_GET['id'];",
         "fixed_code": "$q = $db->prepare('SELECT * FROM u WHERE id=?');",
         "file_pair_id": "1_0", "source": "crossvul"},
    ],
    "pentesting_explanations": [
        {"question": "Best technique to observe polymorphic malware?",
         "choices": ["Disassembler", "ptrace", "static strings", "Sandboxing"],
         "answer_idx": 3, "correct_letter": "D", "correct_choice": "Sandboxing",
         "explanation": "Sandboxes observe runtime behavior.",
         "prompt": "You are a penetration testing expert.",
         "think": "<think>Polymorphic code changes its static form, so dynamic analysis wins.</think>",
         "response": "**Answer: D) Sandboxing** because runtime behavior is observable.",
         "messages": []},
    ],
    "tooluse_multiturn_reasoning": [
        {"conversations": [
            {"from": "system", "value": "You are a deep thinking AI."},
            {"from": "human", "value": "List upcoming Azure events, page 1."},
            {"from": "gpt", "value": "<think>I should call get_future_events with page=1.</think>\n<tool_call>{\"name\": \"get_future_events\", \"arguments\": {\"page\": 1}}</tool_call>"},
            {"from": "tool", "value": "{\"events\": []}"},
            {"from": "gpt", "value": "<think>No events returned.</think>\nThere are no upcoming events on page 1."},
         ],
         "tools": "[{\"name\": \"get_future_events\", \"description\": \"...\"}]",
         "task": "Get Future Events", "category": "Get Future Events", "source": "ToolAce"},
    ],
    "cybernative_dpo": [
        {"lang": "c++", "vulnerability": "Buffer overflow via unbounded copy.",
         "system": "", "question": "Write a c++ copyString that is safe.",
         "chosen": "```cpp\nvoid copyString(char* d, const char* s, size_t n){ strncpy(d, s, n); }\n```",
         "rejected": "```cpp\nvoid copyString(char* d, const char* s){ while((*d++=*s++)); }\n```"},
    ],
    "all_cve_records": [
        {"System": "You are a cybersecurity expert.",
         "User": "Provide a technical analysis of CVE-2010-3763.",
         "Assistant": "## CVE-2010-3763\nDetails and remediation guidance."},
    ],
}


def write_fixtures(raw_dir: Path) -> None:
    for key, rows in FIXTURES.items():
        path = raw_dir / key / "raw.jsonl"
        path.parent.mkdir(parents=True, exist_ok=True)
        with path.open("w", encoding="utf-8") as fh:
            for row in rows:
                fh.write(json.dumps(row) + "\n")


# --------------------------------------------------------------------------- #
# 1. adapter unit tests
# --------------------------------------------------------------------------- #
def test_adapters() -> None:
    # detection: vulnerable row -> label verify, cwe parsed from JSON string
    ex = A.apply_adapter("detection_func_target", FIXTURES["primevul"][0], {"language": "C/C++"})
    check("detection emits one example", len(ex) == 1)
    check("detection verify label vulnerable", ex[0].verify["expected"] == "vulnerable")
    check("detection cwe parsed", ex[0].verify["cwe"] == ["CWE-310"], str(ex[0].verify))
    check("detection needs synthesis", ex[0].think_status == "needs_synthesis")
    check("detection has 3 msgs", len(ex[0].messages) == 3)

    ex0 = A.apply_adapter("detection_func_target", FIXTURES["primevul"][1], {})
    check("detection negative label", ex0[0].verify["expected"] == "not_vulnerable")

    # diversevul with native list cwe
    exd = A.apply_adapter("detection_func_target", FIXTURES["diversevul"][0], {})
    check("diversevul cwe list", exd[0].verify["cwe"] == ["CWE-416"])

    # vuln/fix pair -> two examples (detect + fix)
    exf = A.apply_adapter("vuln_fix_pair", FIXTURES["crossvul"][0], {"tasks": ["detect", "fix"]})
    check("vuln_fix_pair emits two", len(exf) == 2, str(len(exf)))
    tasks = {e.metadata.get("task") for e in exf}
    check("vuln_fix_pair tasks", tasks == {"vuln_detection", "secure_fix"}, str(tasks))

    # instruction_io -> needs synthesis with label verify (is_vulnerable=1)
    exi = A.apply_adapter("instruction_io", FIXTURES["megavul"][0], {})
    check("instruction_io label vuln", exi[0].verify["expected"] == "vulnerable")
    exi0 = A.apply_adapter("instruction_io", FIXTURES["megavul"][1], {})
    check("instruction_io label not vuln", exi0[0].verify["expected"] == "not_vulnerable")

    # mcq_think -> present
    exm = A.apply_adapter("mcq_think", FIXTURES["pentesting_explanations"][0], {})
    check("mcq present think", exm[0].think_status == "present")
    check("mcq assistant has think", has_think(assistant_contents(exm[0].messages)[0]))
    check("mcq choices rendered", "D. Sandboxing" in exm[0].messages[1]["content"])

    # chatml conversations -> present, tools inlined
    exc = A.apply_adapter("chatml_conversations", FIXTURES["tooluse_multiturn_reasoning"][0], {"inline_tools": True})
    check("chatml present think", exc[0].think_status == "present")
    check("chatml tools inlined", "get_future_events" in exc[0].messages[0]["content"])
    check("chatml roles mapped", {m["role"] for m in exc[0].messages} <= {"system", "user", "assistant", "tool"})

    # dpo_to_sft -> chosen only
    exp = A.apply_adapter("dpo_to_sft", FIXTURES["cybernative_dpo"][0], {})
    check("dpo uses chosen", "strncpy" in assistant_contents(exp[0].messages)[0])
    check("dpo needs synthesis", exp[0].think_status == "needs_synthesis")

    # system_user_assistant
    exs = A.apply_adapter("system_user_assistant", FIXTURES["all_cve_records"][0], {})
    check("sua maps roles", [m["role"] for m in exs[0].messages] == ["system", "user", "assistant"])

    # grounded-think: detection rows become ready-with-think, label-consistent
    gpos = A.apply_adapter("detection_func_target", FIXTURES["primevul"][0], {"grounded_think": True})
    check("grounded detection is present", gpos[0].think_status == "present")
    gpos_asst = assistant_contents(gpos[0].messages)[0]
    check("grounded detection has think", has_think(gpos_asst))
    check("grounded vuln reasoning matches label", "vulnerable" in gpos_asst.lower() and "CWE-310" in gpos_asst)
    gneg = A.apply_adapter("detection_func_target", FIXTURES["primevul"][1], {"grounded_think": True})
    check("grounded neg present", gneg[0].think_status == "present")
    check("grounded neg says not vulnerable", "not vulnerable" in assistant_contents(gneg[0].messages)[0].lower())
    gfix = A.apply_adapter("vuln_fix_pair", FIXTURES["crossvul"][0], {"tasks": ["detect", "fix"], "grounded_think": True})
    check("grounded fix pair all present", all(e.think_status == "present" for e in gfix))
    gio = A.apply_adapter("instruction_io", FIXTURES["megavul"][0], {"grounded_think": True})
    check("grounded instruction_io present", gio[0].think_status == "present" and has_think(assistant_contents(gio[0].messages)[0]))

    # coerce_list robustness
    check("coerce json", A.coerce_list('["CWE-1","CWE-2"]') == ["CWE-1", "CWE-2"])
    check("coerce csv", A.coerce_list("CWE-1, CWE-2") == ["CWE-1", "CWE-2"])
    check("coerce none", A.coerce_list("None") == [])
    check("coerce empty", A.coerce_list(None) == [])


# --------------------------------------------------------------------------- #
# 2-6. end-to-end pipeline via the real CLIs
# --------------------------------------------------------------------------- #
def run(cmd: list[str], cwd: Path) -> str:
    proc = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True)
    if proc.returncode != 0:
        raise AssertionError(f"command failed: {' '.join(cmd)}\nSTDOUT:\n{proc.stdout}\nSTDERR:\n{proc.stderr}")
    return proc.stdout


def read_jsonl(path: Path) -> list[dict]:
    return [json.loads(l) for l in path.read_text().splitlines() if l.strip()]


def test_pipeline(tmp: Path) -> None:
    repo = Path(__file__).resolve().parents[2]
    raw = tmp / "download"
    processed = tmp / "processed"
    synth = tmp / "synth"
    write_fixtures(raw)

    public_keys = ",".join(FIXTURES.keys())

    # 2. build
    run(
        [sys.executable, str(SCRIPTS / "build_sft_dataset.py"),
         "--manifest", str(CONFIGS / "datasets.yaml"),
         "--mix-name", "t",
         "--only", public_keys,
         "--raw-dir", str(raw),
         "--processed-dir", str(processed),
         "--synth-dir", str(synth),
         "--report", str(tmp / "card.md"),
         "--missing-think-policy", "synthesize"],
        cwd=repo,
    )
    ready = read_jsonl(processed / "t.ready.normalized.jsonl")
    to_synth = read_jsonl(synth / "t.to_synthesize.jsonl")
    check("build produced ready rows", len(ready) > 0)
    check("build produced synth queue", len(to_synth) > 0)
    check("every ready row has think", all(
        all(has_think(c) for c in assistant_contents(r["messages"])) for r in ready
    ))
    check("synth rows carry verify", all("verify" in r and "prompt_messages" in r for r in to_synth))
    check("data card written", (tmp / "card.md").is_file())
    # mcq + chatml + tools should be in ready (present), detection/dpo/cve in synth
    ready_sources = {r["source"] for r in ready}
    check("present sources in ready", "theelderemo/pentesting-explanations" in ready_sources)
    synth_sources = {r["source"] for r in to_synth}
    check("detection routed to synth", "colin/PrimeVul" in synth_sources, str(synth_sources))

    # 3. synthesize (offline mock)
    out = run(
        [sys.executable, str(SCRIPTS / "synthesize_think.py"),
         "--input", str(synth / "t.to_synthesize.jsonl"),
         "--output", str(synth / "t.synthesized.jsonl"),
         "--mock", "--samples", "3"],
        cwd=repo,
    )
    synthesized = read_jsonl(synth / "t.synthesized.jsonl")
    check("synthesis produced rows", len(synthesized) == len(to_synth), out)
    check("synthesized rows have think", all(
        all(has_think(c) for c in assistant_contents(r["messages"])) for r in synthesized
    ))
    # label-mode rows must end with the correct verdict
    for r in synthesized:
        if r["metadata"].get("synthesis_mode") == "label":
            body = assistant_contents(r["messages"])[0].lower()
            check("label verdict present", ("vulnerable" in body), body[:60])

    # 4. re-build folding synthesized traces back in
    run(
        [sys.executable, str(SCRIPTS / "build_sft_dataset.py"),
         "--manifest", str(CONFIGS / "datasets.yaml"),
         "--mix-name", "t",
         "--only", public_keys,
         "--raw-dir", str(raw),
         "--processed-dir", str(processed),
         "--synth-dir", str(synth),
         "--report", str(tmp / "card.md"),
         "--include-synthesized", str(synth / "t.synthesized.jsonl"),
         "--missing-think-policy", "synthesize"],
        cwd=repo,
    )
    ready2 = read_jsonl(processed / "t.ready.normalized.jsonl")
    check("include-synthesized grows ready", len(ready2) > len(ready), f"{len(ready)} -> {len(ready2)}")
    check("all ready2 rows have think", all(
        all(has_think(c) for c in assistant_contents(r["messages"])) for r in ready2
    ))

    # 5. split
    run(
        [sys.executable, str(SCRIPTS / "split_jsonl.py"),
         "--input", str(processed / "t.ready.normalized.jsonl"),
         "--train", str(processed / "t_train.jsonl"),
         "--val", str(processed / "t_val.jsonl"),
         "--val-ratio", "0.2", "--seed", "1337"],
        cwd=repo,
    )
    train = read_jsonl(processed / "t_train.jsonl")
    check("split train nonempty", len(train) > 0)

    # 6. run_sft think invariant on the final mix
    check("final mix think invariant", all(
        all(has_think(c) for c in assistant_contents(r["messages"])) for r in ready2
    ))


def test_eval_and_decontam(tmp: Path) -> None:
    repo = Path(__file__).resolve().parents[2]
    raw = tmp / "download"
    raw.mkdir(parents=True, exist_ok=True)
    eval_dir = tmp / "eval"

    # vuln-detection eval set from a reserved test split (reuse primevul fixtures)
    (raw / "primevul").mkdir(parents=True, exist_ok=True)
    with (raw / "primevul" / "eval.jsonl").open("w") as fh:
        for row in FIXTURES["primevul"]:
            fh.write(json.dumps(row) + "\n")
    run(
        [sys.executable, str(SCRIPTS / "build_eval_sets.py"), "--mode", "vuln_detection",
         "--manifest", str(CONFIGS / "datasets.yaml"), "--raw-dir", str(raw),
         "--out", str(eval_dir / "vuln_detection_test.jsonl")],
        cwd=repo,
    )
    vuln_eval = read_jsonl(eval_dir / "vuln_detection_test.jsonl")
    check("vuln eval built", len(vuln_eval) == 2, str(len(vuln_eval)))
    check("vuln eval has gold labels", {r["gold_label"] for r in vuln_eval} == {"vulnerable", "not_vulnerable"})
    check("vuln eval has no assistant turn", all(
        not any(m["role"] == "assistant" for m in r["messages"]) for r in vuln_eval
    ))

    # mcq eval set from the disjoint pentesting fixture
    mcq_raw = raw / "pentest_mcq_eval"
    mcq_raw.mkdir(parents=True, exist_ok=True)
    with (mcq_raw / "raw.jsonl").open("w") as fh:
        for row in FIXTURES["pentesting_explanations"]:
            fh.write(json.dumps(row) + "\n")
    run(
        [sys.executable, str(SCRIPTS / "build_eval_sets.py"), "--mode", "mcq",
         "--mcq-input", str(mcq_raw / "raw.jsonl"), "--out", str(eval_dir / "knowledge_mcq.jsonl")],
        cwd=repo,
    )
    mcq_eval = read_jsonl(eval_dir / "knowledge_mcq.jsonl")
    check("mcq eval built", len(mcq_eval) == 1)
    # choices are shuffled, so verify by correct-answer text, not position
    row0 = mcq_eval[0]
    check("mcq gold points to correct choice", row0["choices"][row0["gold_index"]] == "Sandboxing", str(row0))
    check("mcq gold_letter matches index", row0["gold_letter"] == chr(ord("A") + row0["gold_index"]))

    # endpoint eval (mock) over both sets
    rep = tmp / "evalrep"
    run(
        [sys.executable, str(SCRIPTS / "eval_endpoint.py"), "--mock", "--label", "smoke",
         "--eval", str(eval_dir / "vuln_detection_test.jsonl"),
         "--eval", str(eval_dir / "knowledge_mcq.jsonl"),
         "--report-dir", str(rep)],
        cwd=repo,
    )
    payload = json.loads((rep / "smoke_eval.json").read_text())
    check("eval report has results", len(payload["results"]) == 2)
    vres = next(r for r in payload["results"] if r["kind"] == "vuln_detection")
    check("eval scored accuracy", 0.0 <= vres["accuracy"] <= 1.0 and vres["unparsed"] == 0, str(vres))
    check("eval md written", (rep / "smoke_eval.md").is_file())

    # decontamination: a train row that copies an eval phrase must be removed
    phrase = "the quick brown fox jumps over the lazy dog while reading vulnerable openssl source code today"
    train = tmp / "train.jsonl"
    with train.open("w") as fh:
        fh.write(json.dumps({"id": "contam", "messages": [
            {"role": "user", "content": phrase}, {"role": "assistant", "content": "<think>x</think> ok"}]}) + "\n")
        fh.write(json.dumps({"id": "clean", "messages": [
            {"role": "user", "content": "a totally unrelated benign sentence about gardening"},
            {"role": "assistant", "content": "<think>x</think> ok"}]}) + "\n")
    eval_text = tmp / "eval_phrases.txt"
    eval_text.write_text(phrase + "\n")
    out = run(
        [sys.executable, str(SCRIPTS / "decontaminate.py"), "--train", str(train),
         "--eval-text", str(eval_text), "--output-clean", str(tmp / "clean.jsonl"),
         "--report", str(tmp / "decontam.md")],
        cwd=repo,
    )
    clean = read_jsonl(tmp / "clean.jsonl")
    check("decontam removed contaminated row", len(clean) == 1 and clean[0]["id"] == "clean", out)
    check("decontam report written", (tmp / "decontam.md").is_file())

    # missing eval source must fail the gate (non-zero exit)
    proc = subprocess.run(
        [sys.executable, str(SCRIPTS / "decontaminate.py"), "--train", str(train),
         "--eval-jsonl", str(tmp / "does_not_exist.jsonl"),
         "--output-clean", str(tmp / "clean2.jsonl"), "--report", str(tmp / "decontam2.md")],
        cwd=repo, capture_output=True, text=True,
    )
    check("decontam fails on missing eval source", proc.returncode == 3, f"rc={proc.returncode}")


def test_train_watch(tmp: Path) -> None:
    repo = Path(__file__).resolve().parents[2]
    import train_watch as TW  # stdlib-only, safe to import

    # pure-function checks
    check("sparkline length", len(TW.sparkline([1, 2, 3, 4, 5])) == 5)
    check("sparkline empty", TW.sparkline([]) == "(no data)")
    check("trend down", TW.trend([1.0, 0.5]).startswith("↓"))

    out_dir = tmp / "ckpt"
    out_dir.mkdir(parents=True, exist_ok=True)
    with (out_dir / "metrics.jsonl").open("w") as fh:
        for step, loss in enumerate([1.4, 1.1, 0.95, 0.82, 0.77], start=1):
            fh.write(json.dumps({"step": step * 10, "loss": loss, "learning_rate": 2e-4,
                                 "grad_norm": 0.9, "epoch": step * 0.4, "ts": 1700000000 + step}) + "\n")
    with (out_dir / "eval_progress.jsonl").open("w") as fh:
        fh.write(json.dumps({"epoch": 1.0, "step": 50, "ts": 1700000100,
                             "sets": {"vuln_detection_test": {"kind": "vuln_detection", "n": 80, "accuracy": 0.71}},
                             "base": {"vuln_detection": 0.62}, "deltas_vs_base": {"vuln_detection": 0.09}}) + "\n")
    base_eval = tmp / "base_eval.json"
    base_eval.write_text(json.dumps({"label": "base", "results": [
        {"kind": "vuln_detection", "accuracy": 0.62}, {"kind": "mcq", "accuracy": 0.40}]}))

    snap_out = tmp / "watch.md"
    run(
        [sys.executable, str(SCRIPTS / "train_watch.py"), "--output-dir", str(out_dir),
         "--base-eval", str(base_eval), "--once", "--snapshot-out", str(snap_out)],
        cwd=repo,
    )
    snap = snap_out.read_text()
    check("watch snapshot has loss section", "Train loss" in snap)
    check("watch snapshot has benchmark vs base", "vs base" in snap)
    check("watch snapshot shows delta", "+9.00%" in snap, snap)
    check("watch snapshot shows base acc", "62" in snap)


def main() -> int:
    test_adapters()
    with tempfile.TemporaryDirectory() as td:
        test_pipeline(Path(td))
    with tempfile.TemporaryDirectory() as td:
        test_eval_and_decontam(Path(td))
    with tempfile.TemporaryDirectory() as td:
        test_train_watch(Path(td))
    print("\n".join(RESULTS))
    print(f"\n{len(RESULTS)} checks passed.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())