Spaces:
Running on Zero
Running on Zero
| """Test repository migration without downloading any model or reading tokens.""" | |
| import unittest | |
| from unittest.mock import Mock, patch | |
| from fastapi.testclient import TestClient | |
| import app as demo | |
| from runtime import ( | |
| DEFAULT_MODEL_REPO, | |
| LEGACY_MODEL_REPO, | |
| InferenceEngine, | |
| configured_model_source, | |
| ) | |
| class ModelSourceTests(unittest.TestCase): | |
| """Check new defaults, old Space variables, explicit pins and local models.""" | |
| def test_default_repository_and_revision(self) -> None: | |
| """Unset and blank configuration selects the owner's current repository.""" | |
| for environment in [{}, {"LLMON_MODEL_REPO": " ", "LLMON_MODEL_REVISION": ""}]: | |
| with self.subTest(environment=environment), patch.dict("os.environ", environment, clear=True): | |
| self.assertEqual(configured_model_source(), (DEFAULT_MODEL_REPO, "main")) | |
| def test_old_repository_does_not_reuse_its_revision(self) -> None: | |
| """An old pinned commit must never be applied to the new repository.""" | |
| environment = {"LLMON_MODEL_REPO": LEGACY_MODEL_REPO, "LLMON_MODEL_REVISION": "old-repo-only-sha"} | |
| with patch.dict("os.environ", environment, clear=True), self.assertLogs("uvicorn.error", level="WARNING"): | |
| self.assertEqual(configured_model_source(), (DEFAULT_MODEL_REPO, "main")) | |
| def test_explicit_new_repository_pin_is_respected(self) -> None: | |
| """Operators can pin the new checkpoint after verifying its revision.""" | |
| for repo in [DEFAULT_MODEL_REPO, "test-org/explicit-checkpoint"]: | |
| environment = {"LLMON_MODEL_REPO": repo, "LLMON_MODEL_REVISION": "a" * 40} | |
| with self.subTest(repo=repo), patch.dict("os.environ", environment, clear=True): | |
| self.assertEqual(configured_model_source(), (repo, "a" * 40)) | |
| def test_prepare_downloads_migrated_source_and_reports_commit(self) -> None: | |
| """Verify the actual downloader call, not just the settings resolver.""" | |
| commit = "b" * 40 | |
| snapshot = f"/test-only/cache/snapshots/{commit}" | |
| model = Mock() | |
| environment = { | |
| "LLMON_MODEL_REPO": LEGACY_MODEL_REPO, | |
| "LLMON_MODEL_REVISION": "old-repo-only-sha", | |
| "LLMON_DEVICE": "cpu", | |
| } | |
| with ( | |
| patch.dict("os.environ", environment, clear=True), | |
| patch("huggingface_hub.snapshot_download", return_value=snapshot) as download, | |
| patch("vendor.attention_extractor.load_model", return_value=(model, Mock(), {})) as load, | |
| self.assertLogs("uvicorn.error", level="WARNING"), | |
| ): | |
| engine = InferenceEngine() | |
| engine.prepare() | |
| self.assertEqual(download.call_args.kwargs["repo_id"], DEFAULT_MODEL_REPO) | |
| self.assertEqual(download.call_args.kwargs["revision"], "main") | |
| self.assertEqual(download.call_count, 1) | |
| load.assert_called_once_with(snapshot) | |
| model.to.assert_called_once_with("cpu") | |
| self.assertEqual(engine.status, "ready") | |
| with patch.object(demo, "engine", engine): | |
| health = TestClient(demo.app).get("/api/health").json() | |
| self.assertEqual(health["model_repo"], DEFAULT_MODEL_REPO) | |
| self.assertEqual(health["model_revision"], commit) | |
| self.assertEqual(health["model_requested_revision"], "main") | |
| self.assertNotIn(snapshot, str(health)) | |
| def test_local_checkpoint_takes_precedence_without_download(self) -> None: | |
| """Keep explicit CPU/MPS development paths separate from Hub models.""" | |
| environment = {"LLMON_MODEL_PATH": "/test-only/local", "LLMON_DEVICE": "cpu"} | |
| with ( | |
| patch.dict("os.environ", environment, clear=True), | |
| patch("huggingface_hub.snapshot_download") as download, | |
| patch("vendor.attention_extractor.load_model", return_value=(Mock(), Mock(), {})) as load, | |
| ): | |
| engine = InferenceEngine() | |
| engine.prepare() | |
| download.assert_not_called() | |
| load.assert_called_once_with("/test-only/local") | |
| self.assertIsNone(engine.model_info["model_repo"]) | |
| self.assertIsNone(engine.model_info["model_revision"]) | |
| def test_failed_download_never_falls_back_to_old_model(self) -> None: | |
| """An unavailable release must surface an error, not a different model.""" | |
| with ( | |
| patch.dict("os.environ", {"LLMON_DEVICE": "cpu"}, clear=True), | |
| patch("huggingface_hub.snapshot_download", side_effect=OSError("unavailable")) as download, | |
| ): | |
| engine = InferenceEngine() | |
| with self.assertRaises(OSError): | |
| engine.prepare() | |
| self.assertEqual(download.call_count, 1) | |
| self.assertEqual(download.call_args.kwargs["repo_id"], DEFAULT_MODEL_REPO) | |
| self.assertIsNone(engine.model) | |
| self.assertNotEqual(engine.status, "ready") | |
| if __name__ == "__main__": | |
| unittest.main() | |