Rebasing and fixing UTs

Author: kumar-shivam-ranjan

ads/aqua/extension/model_handler.py

@@ -10,15 +10,14 @@
 from ads.aqua.common.decorator import handle_exceptions
 from ads.aqua.common.errors import AquaRuntimeError, AquaValueError
 from ads.aqua.common.utils import (
-    get_container_config,
     get_hf_model_info,
     list_hf_models,
 )
 from ads.aqua.extension.base_handler import AquaAPIhandler
 from ads.aqua.extension.errors import Errors
 from ads.aqua.model import AquaModelApp
 from ads.aqua.model.entities import AquaModelSummary, HFModelSummary
-from ads.aqua.ui import AquaContainerConfig, ModelFormat
+from ads.aqua.ui import ModelFormat
 
 
 class AquaModelHandler(AquaAPIhandler):
@@ -154,14 +153,11 @@ def put(self, id):
             raise HTTPError(400, Errors.NO_INPUT_DATA)
 
         inference_container = input_data.get("inference_container")
-        containers = list(
-            AquaContainerConfig.from_container_index_json(
-                config=get_container_config(), enable_spec=True
-            ).inference.values()
-        )
-        family_values = [item.family for item in containers]
-
-        if inference_container is not None and inference_container not in family_values:
+        inference_containers = AquaModelApp.list_valid_inference_containers()
+        if (
+            inference_container is not None
+            and inference_container not in inference_containers
+        ):
             raise HTTPError(
                 400, Errors.INVALID_VALUE_OF_PARAMETER.format("inference_container")
             )

ads/aqua/model/model.py

@@ -26,6 +26,7 @@
     copy_model_config,
     create_word_icon,
     get_artifact_path,
+    get_container_config,
     get_hf_model_info,
     list_os_files_with_extension,
     load_config,
@@ -718,6 +719,16 @@ def clear_model_list_cache(
                 }
         return res
 
+    @staticmethod
+    def list_valid_inference_containers():
+        containers = list(
+            AquaContainerConfig.from_container_index_json(
+                config=get_container_config(), enable_spec=True
+            ).inference.values()
+        )
+        family_values = [item.family for item in containers]
+        return family_values
+
     def _create_model_catalog_entry(
         self,
         os_path: str,

tests/unitary/with_extras/aqua/test_model_handler.py

@@ -20,7 +20,6 @@
 )
 from ads.aqua.model import AquaModelApp
 from ads.aqua.model.entities import AquaModel, AquaModelSummary, HFModelSummary
-from ads.aqua.ui import AquaContainerConfig
 
 
 class ModelHandlerTestCase(TestCase):
@@ -94,18 +93,15 @@ def test_delete_with_id(self, mock_delete, mock_urlparse):
             mock_urlparse.assert_called()
             mock_delete.assert_called()
 
-    @patch.object(AquaContainerConfig, "from_container_index_json")
+    @patch.object(AquaModelApp, "list_valid_inference_containers")
     @patch.object(AquaModelApp, "edit_registered_model")
-    def test_put(self, mock_edit, mock_container_index):
+    def test_put(self, mock_edit, mock_inference_container_list):
         mock_edit.return_value = {"state": "EDITED"}
-        mock_inference = MagicMock()
-        mock_inference.values.return_value = [
-            MagicMock(family="odsc-vllm-serving"),
-            MagicMock(family="odsc-tgi-serving"),
-            MagicMock(family="odsc-vllm-serving"),
+        mock_inference_container_list.return_value = [
+            "odsc-vllm-serving",
+            "odsc-tgi-serving",
+            "odsc-llama-cpp-serving",
         ]
-
-        mock_container_index.return_value = MagicMock(inference=mock_inference)
         self.model_handler.get_json_body = MagicMock(
             return_value=dict(
                 task="text_generation",
@@ -118,9 +114,9 @@ def test_put(self, mock_edit, mock_container_index):
         ) as mock_finish:
             mock_finish.side_effect = lambda x: x
             result = self.model_handler.put(id="ocid1.datasciencemodel.oc1.iad.xxx")
-            print(f"result: ", result)
             assert result["state"] is "EDITED"
-            mock_edit.assert_called()
+            mock_edit.assert_called_once()
+            mock_inference_container_list.assert_called_once()
 
     @patch.object(AquaModelApp, "list")
     def test_list(self, mock_list):