tests/test_registry.py

from pathlib import Path
import tempfile
import unittest
from baim.registry import Registry,file_hash


class RegistryTests(unittest.TestCase):
    def test_gates_integrity_and_rollback(self):
        with tempfile.TemporaryDirectory() as temp:
            root = Path(temp)
            source = root/'checkpoint'
            source.mkdir()
            (source/'model.safetensors').write_bytes(b'unit-test-fixture-not-a-model')
            registry = Registry(root/'registry')
            try:
                registry.register('v000',source,{})
                with self.assertRaisesRegex(ValueError,'MISSING_EVIDENCE'):
                    registry.promote('v000')
                self.assertIsNone(registry.champion)
                metrics = dict(suite_sha256='a'*64,artifact_sha256=file_hash(source/'model.safetensors'),
                    task_success_rate=.9,out_of_domain_success=.9,recovery_rate=.9,unsafe_action_rate=0,
                    premature_finish_rate=0,loop_rate=0,p95_action_ms=1000,peak_service_rss_bytes=1024**3,
                    target_hardware='linux-epyc9354p-2vcpu',evaluator_version='unit-test-only')
                registry.register('v001',source,metrics,parent='v000')
                registry.promote('v001')
                self.assertEqual(registry.champion,'v001')
                registry.register('v002',source,{**metrics,'p95_action_ms':900},parent='v001')
                registry.promote('v002')
                registry.rollback('v001')
                self.assertEqual(registry.champion,'v001')
                (root/'registry/v002/model.safetensors').write_bytes(b'changed')
                with self.assertRaisesRegex(ValueError,'integrity'):
                    registry.promote('v002')
                self.assertEqual(registry.champion,'v001')
                with self.assertRaisesRegex(ValueError,'immutable'):
                    registry.register('v001',source,metrics)
            finally:
                registry.close()