Files
aituner/runs/frontier-code-trace-v0/test_code_trace_preflight.py

81 lines
2.8 KiB
Python

#!/usr/bin/env python3
from __future__ import annotations
import json
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path
ROOT = Path(__file__).resolve().parent
class CodeTracePreflightTest(unittest.TestCase):
def test_audit_and_materialize_512_block_window(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
root = Path(tmp)
source = root / "051315-051317.jsonl"
with source.open("w") as stream:
for index in range(4501):
input_tokens = 513 if index % 2 else 512
stream.write(
json.dumps(
{
"chat_id": index,
"parent_chat_id": index - 1 if index % 2 else -1,
"timestamp": float(index),
"input_length": input_tokens,
"output_length": 32,
"hash_ids": [index // 2]
if input_tokens == 512
else [index // 2, 100000 + index],
}
)
+ "\n"
)
audit = root / "audit.json"
subprocess.run(
[
sys.executable,
str(ROOT / "audit_code_trace.py"),
"--source",
str(source),
"--output",
str(audit),
],
check=True,
)
payload = json.loads(audit.read_text())
self.assertEqual(payload["data_gate"], "PASS")
self.assertEqual(
payload["selected"]["hash_contract"]["exact_source_block_size"], 512
)
self.assertEqual(payload["max_model_len_recommendation"], 40960)
output = root / "window"
subprocess.run(
[
sys.executable,
str(ROOT / "prepare_code_window.py"),
"--audit",
str(audit),
"--output-root",
str(output),
],
check=True,
)
manifest = json.loads((output / "window-manifest.json").read_text())
rows = [
json.loads(line)
for line in (output / "code-raw-window.jsonl").read_text().splitlines()
]
self.assertEqual(manifest["requests"], len(rows))
self.assertGreaterEqual(manifest["duration_s"], 3600)
self.assertEqual(rows[0]["sampling_u"], rows[1]["sampling_u"])
if __name__ == "__main__":
unittest.main()