test: fix test_registry

jjleng · Apr 19, 2024 · 539f453 · 539f453
1 parent be3f8e4
commit 539f453
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/tests/model/test_registry.py b/tests/model/test_registry.py
@@ -1,9 +1,14 @@
+from unittest.mock import MagicMock, Mock, patch
+
 import pytest
 
-from paka.model.registry import HuggingFaceModel, get_model, model_registry
+from paka.model.hf_model import BaseMLModel
+
 
+@patch.object(BaseMLModel, "get_model_store", return_value=MagicMock())
+def test_get_model(get_model_store_mock: Mock) -> None:
+    from paka.model.registry import HuggingFaceModel, get_model, model_registry
 
-def test_get_model() -> None:
     model_registry.append(
         HuggingFaceModel(
             name="test_model",