Skip to content

Commit 240146b

Browse files
committed
add test
1 parent 9a73598 commit 240146b

File tree

1 file changed

+30
-0
lines changed

1 file changed

+30
-0
lines changed

tests/unit/sagemaker/jumpstart/hub/test_hub.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -192,6 +192,36 @@ def test_describe_model_success(mock_describe_hub_content_response, sagemaker_se
192192
)
193193

194194

195+
@patch("sagemaker.jumpstart.hub.interfaces.DescribeHubContentResponse.from_json")
196+
def test_describe_model_one_thrown_error(mock_describe_hub_content_response, sagemaker_session):
197+
mock_describe_hub_content_response.return_value = Mock()
198+
mock_list_hub_content_versions = sagemaker_session.list_hub_content_versions
199+
mock_list_hub_content_versions.return_value = {
200+
"HubContentSummaries": [
201+
{"HubContentVersion": "1.0"},
202+
{"HubContentVersion": "2.0"},
203+
{"HubContentVersion": "3.0"},
204+
]
205+
}
206+
mock_describe_hub_content = sagemaker_session.describe_hub_content
207+
mock_describe_hub_content.side_effect = [
208+
Exception("Some exception"),
209+
{"HubContentName": "test-model", "HubContentVersion": "3.0"},
210+
]
211+
212+
hub = Hub(hub_name=HUB_NAME, sagemaker_session=sagemaker_session)
213+
214+
with patch("sagemaker.jumpstart.hub.utils.get_hub_model_version") as mock_get_hub_model_version:
215+
mock_get_hub_model_version.return_value = "3.0"
216+
217+
hub.describe_model("test-model")
218+
219+
mock_describe_hub_content.asssert_called_times(2)
220+
mock_describe_hub_content.assert_called_with(
221+
hub_name=HUB_NAME, hub_content_name="test-model", hub_content_version="3.0", hub_content_type="Model"
222+
)
223+
224+
195225
def test_create_hub_content_reference(sagemaker_session):
196226
hub = Hub(hub_name=HUB_NAME, sagemaker_session=sagemaker_session)
197227
model_name = "mock-model-one-huggingface"

0 commit comments

Comments
 (0)