tests/test_model_loader.py

import unittest


class ModelManagerTests(unittest.TestCase):
    def test_manager_caches_same_key_and_replaces_changed_key(self):
        from model_loader import ModelManager

        calls = []
        manager = ModelManager(
            loader=lambda model, path, config: calls.append((model, path, config)) or object()
        )
        first = manager.get("owner/model", "a", "pi05_ur_demo_no_state")
        self.assertIs(manager.get("owner/model", "a", "pi05_ur_demo_no_state"), first)
        second = manager.get("owner/model", "a", "pi05_ur_demo_state")
        self.assertIsNot(second, first)
        self.assertEqual(calls, [
            ("owner/model", "a", "pi05_ur_demo_no_state"),
            ("owner/model", "a", "pi05_ur_demo_state"),
        ])

    def test_failed_load_is_reported_and_can_retry(self):
        from model_loader import ModelManager, ModelUnavailableError

        attempts = 0

        def loader(*_):
            nonlocal attempts
            attempts += 1
            raise RuntimeError("bad checkpoint")

        manager = ModelManager(loader=loader)
        for _ in range(2):
            with self.assertRaisesRegex(ModelUnavailableError, "bad checkpoint"):
                manager.get("owner/model", "a", "pi05_ur_demo_no_state")
        self.assertEqual(attempts, 2)
        self.assertIn("bad checkpoint", manager.health_message)

    def test_invalid_identity_is_rejected_before_loader(self):
        from model_loader import ModelManager

        manager = ModelManager(loader=lambda *_: self.fail("loader was called"))
        with self.assertRaises(ValueError):
            manager.get("", "checkpoint", "pi05_ur_demo_no_state")

    def test_unsupported_policy_config_is_rejected_before_loader(self):
        from model_loader import ModelManager

        manager = ModelManager(loader=lambda *_: self.fail("loader was called"))
        with self.assertRaisesRegex(ValueError, "unsupported policy config"):
            manager.get("owner/model", "checkpoint", "pi05_libero")


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