nucleic-purpose-classifier/tests/test_prepare_data.py

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()