| import copy |
| import json |
| from pathlib import Path |
| import tempfile |
| import unittest |
|
|
| from data_transition import (SPLITS, PROTECTED_SPLITS, data_signature, |
| sha256, verify_train_data_transition) |
|
|
|
|
| class DataTransitionTests(unittest.TestCase): |
| def setUp(self): |
| self.directory = tempfile.TemporaryDirectory() |
| self.addCleanup(self.directory.cleanup) |
| root = Path(self.directory.name) |
| self.old, self.new = root / "old", root / "new" |
| self.provenance = {"revision": "pinned-model"} |
| for path in (self.old, self.new): |
| path.mkdir() |
| for split in SPLITS: |
| (path / f"{split}.jsonl").write_text('{"id":"original"}\n') |
| (self.new / "train.jsonl").write_text('{"id":"original"}\n{"id":"new"}\n') |
| self.parent = {"split_sha256": self.hashes(self.old)} |
| (self.old / "manifest.json").write_text(json.dumps(self.parent)) |
| self.current = {"split_sha256": self.hashes(self.new), "base_dataset": { |
| "path": str(self.old), "manifest_sha256": sha256(self.old / "manifest.json"), |
| "split_sha256": self.parent["split_sha256"]}} |
| self.write_current() |
| self.saved = {"config": {"dataset": str(self.old), "max_tokens": 512}, |
| "model_provenance": self.provenance, |
| "data_signature": data_signature(self.parent, self.provenance, 512)} |
| self.config = {"dataset": str(self.new), "max_tokens": 512} |
|
|
| def hashes(self, path): |
| return {split: sha256(path / f"{split}.jsonl") for split in SPLITS} |
|
|
| def write_current(self): |
| (self.new / "manifest.json").write_text(json.dumps(self.current)) |
|
|
| def verify(self): |
| return verify_train_data_transition(self.saved, self.config, |
| data_signature(self.current, self.provenance, 512)) |
|
|
| def test_expansion_records_both_signatures_and_protected_hashes(self): |
| untouched = copy.deepcopy(self.saved) |
| result = self.verify() |
| self.assertEqual(self.saved, untouched) |
| self.assertEqual(result["parent_data_signature"], self.saved["data_signature"]) |
| self.assertNotEqual(result["parent_data_signature"], result["data_signature"]) |
| self.assertEqual(set(result["protected_split_sha256"]), set(PROTECTED_SPLITS)) |
|
|
| def test_reserved_changes_rejected_even_with_valid_updated_hash(self): |
| for split in PROTECTED_SPLITS: |
| with self.subTest(split=split): |
| original = (self.new / f"{split}.jsonl").read_bytes() |
| (self.new / f"{split}.jsonl").write_text('changed') |
| self.current["split_sha256"] = self.hashes(self.new) |
| self.write_current() |
| with self.assertRaisesRegex(ValueError, "reserved"): |
| self.verify() |
| (self.new / f"{split}.jsonl").write_bytes(original) |
|
|
| def test_actual_bytes_missing_splits_and_false_lineage_rejected(self): |
| (self.new / "train.jsonl").write_text('corrupted') |
| with self.assertRaisesRegex(ValueError, "checksum"): |
| self.verify() |
| self.current["split_sha256"] = self.hashes(self.new) |
| self.current["base_dataset"]["manifest_sha256"] = "wrong" |
| self.write_current() |
| with self.assertRaisesRegex(ValueError, "lineage"): |
| self.verify() |
| del self.current["split_sha256"]["holdout"] |
| self.write_current() |
| with self.assertRaisesRegex(ValueError, "five registered"): |
| self.verify() |
|
|
| def test_wrong_parent_signature_or_token_policy_rejected(self): |
| self.saved["data_signature"] = "incorrect" |
| with self.assertRaisesRegex(ValueError, "Parent data signature"): |
| self.verify() |
| self.config["max_tokens"] = 1024 |
| with self.assertRaisesRegex(ValueError, "max_tokens"): |
| self.verify() |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
