Skip to content

Commit ad0ce2e

Browse files
committed
Added uri UTs.
1 parent 9f69e35 commit ad0ce2e

File tree

1 file changed

+72
-2
lines changed

1 file changed

+72
-2
lines changed

tests/unitary/default_setup/model/test_datascience_model.py

+72-2
Original file line numberDiff line numberDiff line change
@@ -1132,6 +1132,57 @@ def test_download_artifact_for_model_created_by_reference(
11321132

11331133
mock_large_download.assert_called()
11341134

1135+
@patch("ads.common.oci_client.OCIClientFactory")
1136+
def test_add_artifact_uri(self, mock_oci_client_factory):
1137+
r = ObjectVersionSummary()
1138+
r.name = "model-linear-2/linear-2.pkl"
1139+
r.size = 566
1140+
r.time_modified = "2024-04-22T12:34:26.670000+00:00"
1141+
r.version_id = "dc26a7d2-8041-4b37-8ed0-9e8c10869340"
1142+
resp = [r]
1143+
1144+
# Mock response object
1145+
mock_response = MagicMock()
1146+
mock_response.data.items = resp
1147+
mock_response.has_next_page = False
1148+
mock_response.next_page = None
1149+
1150+
# Mock object storage client
1151+
mock_object_storage_client = MagicMock()
1152+
mock_object_storage_client.list_object_versions.return_value = mock_response
1153+
1154+
mock_oci_client_factory.return_value.object_storage = mock_object_storage_client
1155+
1156+
# self.mock_dsc_model
1157+
self.mock_dsc_model.add_artifact(uri="oci://bucket@namespace/prefix")
1158+
expected_out = {
1159+
"version": "1.0",
1160+
"type": "modelOSSReferenceDescription",
1161+
"models": [
1162+
{
1163+
"namespace": "namespace",
1164+
"bucketName": "bucket",
1165+
"prefix": "prefix",
1166+
"objects": [
1167+
{
1168+
"name": "model-linear-2/linear-2.pkl",
1169+
"version": "dc26a7d2-8041-4b37-8ed0-9e8c10869340",
1170+
"sizeInBytes": 566,
1171+
}
1172+
],
1173+
}
1174+
],
1175+
}
1176+
assert self.mock_dsc_model.model_file_description == expected_out
1177+
self.mock_dsc_model.remove_artifact(uri="oci://bucket@namespace/prefix")
1178+
assert self.mock_dsc_model.model_file_description != expected_out
1179+
expected_out = {
1180+
"version": "1.0",
1181+
"type": "modelOSSReferenceDescription",
1182+
"models": [],
1183+
}
1184+
assert self.mock_dsc_model.model_file_description == expected_out
1185+
11351186
@patch("ads.common.oci_client.OCIClientFactory")
11361187
def test_add_artifact(self, mock_oci_client_factory):
11371188
r = ObjectVersionSummary()
@@ -1189,8 +1240,28 @@ def test_add_artifact(self, mock_oci_client_factory):
11891240
}
11901241
assert self.mock_dsc_model.model_file_description == expected_out
11911242

1243+
def test_remove_artifact_uri(self):
1244+
self.mock_dsc_model.remove_artifact(
1245+
uri="oci://unzip-multi-model@ociodscdev/model-linear-1"
1246+
)
1247+
1248+
assert self.mock_dsc_model.model_file_description == None
1249+
1250+
self.mock_dsc_model.set_spec(
1251+
CONST_MODEL_FILE_DESCRIPTION, deepcopy(MODEL_BY_REF_JSON)
1252+
)
1253+
assert self.mock_dsc_model.model_file_description == MODEL_BY_REF_JSON
1254+
1255+
self.mock_dsc_model.remove_artifact(
1256+
uri="oci://unzip-multi-model@ociodscdev/model-linear-1"
1257+
)
1258+
assert self.mock_dsc_model.model_file_description != MODEL_BY_REF_JSON
1259+
1260+
exptected_json = deepcopy(MODEL_BY_REF_JSON)
1261+
exptected_json["models"] = exptected_json["models"][1:]
1262+
assert self.mock_dsc_model.model_file_description == exptected_json
1263+
11921264
def test_remove_artifact(self):
1193-
# self.mock_dsc_model.remove_artifact(uri="oci://unzip-multi-model@ociodscdev/model-linear-1")
11941265
self.mock_dsc_model.remove_artifact(
11951266
namespace="ociodscdev", bucket="unzip-multi-model", prefix="model-linear-1"
11961267
)
@@ -1202,7 +1273,6 @@ def test_remove_artifact(self):
12021273
)
12031274
assert self.mock_dsc_model.model_file_description == MODEL_BY_REF_JSON
12041275

1205-
# self.mock_dsc_model.remove_artifact(uri="oci://unzip-multi-model@ociodscdev/model-linear-1")
12061276
self.mock_dsc_model.remove_artifact(
12071277
namespace="ociodscdev", bucket="unzip-multi-model", prefix="model-linear-1"
12081278
)

0 commit comments

Comments
 (0)