infosec-v1 / code /training /tests /test_data_pipeline.py

Super-squash branch 'main' using huggingface_hub

994182c 4 months ago

21.5 kB

#!/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())

	#!/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())