nucleic-purpose-classifier/tests/test_prepare_history_experiment.py

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_history_experiment
import purpose_data


def example(index: int, **overrides):
    value = {
        "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": "core",
        "lang": "en",
    }
    value.update(overrides)
    return value


class PrepareHistoryExperimentTests(unittest.TestCase):
    def test_history_is_training_only_and_eval_splits_are_byte_identical(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            base = root / "base"
            base.mkdir()
            train = [example(0)]
            validation = [example(1)]
            test = [example(2)]
            purpose_data.write_jsonl(base / "train.jsonl", train)
            purpose_data.write_jsonl(base / "validation.jsonl", validation)
            purpose_data.write_jsonl(base / "test.jsonl", test)
            validation_bytes = (base / "validation.jsonl").read_bytes()
            test_bytes = (base / "test.jsonl").read_bytes()

            history_path = root / "history.jsonl"
            history = [
                example(10, prompt="Add a durable upload endpoint", purpose="backendImpl"),
                example(11, prompt=validation[0]["prompt"], purpose=validation[0]["purpose"]),
                example(12, prompt="Can you assess this?", slice="vague-eval"),
            ]
            purpose_data.write_jsonl(history_path, history)
            fixtures = root / "fixtures.json"
            fixtures.write_text(
                json.dumps(
                    [{"prompt": "Review the release diff", "purpose": "review"}]
                ),
                encoding="utf-8",
            )
            output = root / "output"

            manifest = prepare_history_experiment.prepare(
                base_dataset=base,
                history_path=history_path,
                fixtures_path=fixtures,
                output_dir=output,
                near_duplicate_threshold=0.92,
                overwrite_output=False,
            )

            self.assertEqual(2, manifest["outputs"]["train"]["records"])
            self.assertEqual(1, manifest["augmentation"]["acceptedHistoryRecords"])
            self.assertEqual(
                {"exact-eval-overlap": 1, "vague-eval": 1},
                manifest["augmentation"]["exclusions"],
            )
            self.assertEqual(validation_bytes, (output / "validation.jsonl").read_bytes())
            self.assertEqual(test_bytes, (output / "test.jsonl").read_bytes())
            self.assertNotIn(
                history[0]["prompt"],
                (output / "validation.jsonl").read_text(encoding="utf-8"),
            )
            self.assertNotIn(
                history[0]["prompt"],
                (output / "test.jsonl").read_text(encoding="utf-8"),
            )

    def test_near_eval_overlap_is_excluded_and_output_fails_closed(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            base = root / "base"
            base.mkdir()
            purpose_data.write_jsonl(base / "train.jsonl", [example(0)])
            words = [f"token{index}" for index in range(100)]
            validation_prompt = " ".join(words)
            purpose_data.write_jsonl(
                base / "validation.jsonl",
                [example(1, prompt=validation_prompt, purpose="review")],
            )
            purpose_data.write_jsonl(base / "test.jsonl", [example(2)])
            words[50] = "replacement"
            history_path = root / "history.jsonl"
            purpose_data.write_jsonl(
                history_path,
                [example(10, prompt=" ".join(words), purpose="review")],
            )
            fixtures = root / "fixtures.json"
            fixtures.write_text("[]", encoding="utf-8")
            output = root / "output"

            manifest = prepare_history_experiment.prepare(
                base_dataset=base,
                history_path=history_path,
                fixtures_path=fixtures,
                output_dir=output,
                near_duplicate_threshold=0.92,
                overwrite_output=False,
            )
            self.assertEqual(
                {"near-eval-overlap": 1},
                manifest["augmentation"]["exclusions"],
            )
            with self.assertRaisesRegex(purpose_data.DataError, "output is not empty"):
                prepare_history_experiment.prepare(
                    base_dataset=base,
                    history_path=history_path,
                    fixtures_path=fixtures,
                    output_dir=output,
                    near_duplicate_threshold=0.92,
                    overwrite_output=False,
                )


if __name__ == "__main__":
    unittest.main()