import csv import json import sys import tempfile import unittest from pathlib import Path MODULE_DIR = Path(__file__).resolve().parents[1] sys.path.insert(0, str(MODULE_DIR)) import prepare_data import purpose_data import review_contract def example(index: int): return { "prompt": f"Implement sample endpoint number {index} with stable pagination", "purpose": purpose_data.LABELS[index % len(purpose_data.LABELS)], "secondary": None, "mixed": False, "difficulty": 0.4, "slice": "vague-eval" if index < 5 else "core", "lang": "en", } class PrepareIntegrationTests(unittest.TestCase): def test_refresh_then_verify_frozen_split(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) source = root / "source.jsonl" purpose_data.write_jsonl(source, (example(index) for index in range(80))) fixtures = root / "fixtures.json" fixtures.write_text( json.dumps( [ {"prompt": "Plan the cache migration", "purpose": "planning"}, {"prompt": "Anything else?", "purpose": "general"}, ] ), encoding="utf-8", ) output = root / "output" frozen = root / "frozen.jsonl" manifest = root / "manifest.json" first = prepare_data.prepare( sources=[source], fixtures_path=fixtures, output_dir=output, frozen_test_path=frozen, manifest_path=manifest, refresh_frozen_test=True, seed=23, near_duplicate_threshold=0.92, ) second = prepare_data.prepare( sources=[source], fixtures_path=fixtures, output_dir=output, frozen_test_path=frozen, manifest_path=manifest, refresh_frozen_test=False, seed=23, near_duplicate_threshold=0.92, ) self.assertEqual(first, second) self.assertEqual(65, first["splits"]["train"]["records"]) self.assertEqual(8, first["splits"]["validation"]["records"]) self.assertEqual(8, first["splits"]["test"]["logicalRecords"]) train = purpose_data.load_jsonl(output / "train.jsonl") self.assertFalse(any(row["slice"] == "vague-eval" for row in train)) def test_completed_human_review_is_applied_and_recorded(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) source = root / "source.jsonl" purpose_data.write_jsonl(source, (example(index) for index in range(80))) fixtures = root / "fixtures.json" fixtures.write_text( json.dumps( [{"prompt": "Plan the cache migration", "purpose": "planning"}] ), encoding="utf-8", ) source_records = purpose_data.load_sources([source]) review_population = purpose_data.curate_records( source_records, purpose_data.load_classifiable_fixtures(fixtures), ).records review_csv = root / "review.csv" review_contract.write_review_csv(review_csv, review_population) with review_csv.open("r", encoding="utf-8", newline="") as handle: rows = list(csv.DictReader(handle)) for row in rows: row["reviewStatus"] = "accept" rows[0]["reviewStatus"] = "reject" rows[0]["reviewNotes"] = "Not classifiable after human review." with review_csv.open("w", encoding="utf-8", newline="") as handle: writer = csv.DictWriter( handle, fieldnames=review_contract.REVIEW_CSV_FIELDS, ) writer.writeheader() writer.writerows(rows) human_review = root / "human-review.json" review_contract.finalize_human_review( review_csv, human_review, review_population, dataset_version=prepare_data.DATASET_VERSION, fraction=1.0, seed=41, ) manifest = prepare_data.prepare( sources=[source], fixtures_path=fixtures, output_dir=root / "output", frozen_test_path=root / "frozen.jsonl", manifest_path=root / "manifest.json", refresh_frozen_test=True, seed=23, near_duplicate_threshold=0.92, human_review_path=human_review, ) self.assertEqual(79, manifest["curation"]["retainedRecords"]) self.assertEqual( 1, manifest["curation"]["humanReview"]["summary"]["rejected"], ) if __name__ == "__main__": unittest.main()