#!/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- 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 "" in content and "" 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": "Polymorphic code changes its static form, so dynamic analysis wins.", "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": "I should call get_future_events with page=1.\n{\"name\": \"get_future_events\", \"arguments\": {\"page\": 1}}"}, {"from": "tool", "value": "{\"events\": []}"}, {"from": "gpt", "value": "No events returned.\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": "x ok"}]}) + "\n") fh.write(json.dumps({"id": "clean", "messages": [ {"role": "user", "content": "a totally unrelated benign sentence about gardening"}, {"role": "assistant", "content": "x 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())