Added uri UTs.

tests/unitary/default_setup/model/test_datascience_model.py

@@ -1132,6 +1132,57 @@ def test_download_artifact_for_model_created_by_reference(
 
         mock_large_download.assert_called()
 
+    @patch("ads.common.oci_client.OCIClientFactory")
+    def test_add_artifact_uri(self, mock_oci_client_factory):
+        r = ObjectVersionSummary()
+        r.name = "model-linear-2/linear-2.pkl"
+        r.size = 566
+        r.time_modified = "2024-04-22T12:34:26.670000+00:00"
+        r.version_id = "dc26a7d2-8041-4b37-8ed0-9e8c10869340"
+        resp = [r]
+
+        # Mock response object
+        mock_response = MagicMock()
+        mock_response.data.items = resp
+        mock_response.has_next_page = False
+        mock_response.next_page = None
+
+        # Mock object storage client
+        mock_object_storage_client = MagicMock()
+        mock_object_storage_client.list_object_versions.return_value = mock_response
+
+        mock_oci_client_factory.return_value.object_storage = mock_object_storage_client
+
+        # self.mock_dsc_model
+        self.mock_dsc_model.add_artifact(uri="oci://bucket@namespace/prefix")
+        expected_out = {
+            "version": "1.0",
+            "type": "modelOSSReferenceDescription",
+            "models": [
+                {
+                    "namespace": "namespace",
+                    "bucketName": "bucket",
+                    "prefix": "prefix",
+                    "objects": [
+                        {
+                            "name": "model-linear-2/linear-2.pkl",
+                            "version": "dc26a7d2-8041-4b37-8ed0-9e8c10869340",
+                            "sizeInBytes": 566,
+                        }
+                    ],
+                }
+            ],
+        }
+        assert self.mock_dsc_model.model_file_description == expected_out
+        self.mock_dsc_model.remove_artifact(uri="oci://bucket@namespace/prefix")
+        assert self.mock_dsc_model.model_file_description != expected_out
+        expected_out = {
+            "version": "1.0",
+            "type": "modelOSSReferenceDescription",
+            "models": [],
+        }
+        assert self.mock_dsc_model.model_file_description == expected_out
+
     @patch("ads.common.oci_client.OCIClientFactory")
     def test_add_artifact(self, mock_oci_client_factory):
         r = ObjectVersionSummary()
@@ -1189,8 +1240,28 @@ def test_add_artifact(self, mock_oci_client_factory):
         }
         assert self.mock_dsc_model.model_file_description == expected_out
 
+    def test_remove_artifact_uri(self):
+        self.mock_dsc_model.remove_artifact(
+            uri="oci://unzip-multi-model@ociodscdev/model-linear-1"
+        )
+
+        assert self.mock_dsc_model.model_file_description == None
+
+        self.mock_dsc_model.set_spec(
+            CONST_MODEL_FILE_DESCRIPTION, deepcopy(MODEL_BY_REF_JSON)
+        )
+        assert self.mock_dsc_model.model_file_description == MODEL_BY_REF_JSON
+
+        self.mock_dsc_model.remove_artifact(
+            uri="oci://unzip-multi-model@ociodscdev/model-linear-1"
+        )
+        assert self.mock_dsc_model.model_file_description != MODEL_BY_REF_JSON
+
+        exptected_json = deepcopy(MODEL_BY_REF_JSON)
+        exptected_json["models"] = exptected_json["models"][1:]
+        assert self.mock_dsc_model.model_file_description == exptected_json
+
     def test_remove_artifact(self):
-        # self.mock_dsc_model.remove_artifact(uri="oci://unzip-multi-model@ociodscdev/model-linear-1")
         self.mock_dsc_model.remove_artifact(
             namespace="ociodscdev", bucket="unzip-multi-model", prefix="model-linear-1"
         )
@@ -1202,7 +1273,6 @@ def test_remove_artifact(self):
         )
         assert self.mock_dsc_model.model_file_description == MODEL_BY_REF_JSON
 
-        # self.mock_dsc_model.remove_artifact(uri="oci://unzip-multi-model@ociodscdev/model-linear-1")
         self.mock_dsc_model.remove_artifact(
             namespace="ociodscdev", bucket="unzip-multi-model", prefix="model-linear-1"
         )