fix: Adding more code coverage

Author: chrstfu

src/sagemaker/jumpstart/hub/utils.py

@@ -219,18 +219,17 @@ def get_hub_model_version(
     except Exception as ex:
         raise Exception(f"Failed calling list_hub_content_versions: {str(ex)}")
 
-    marketplace_hub_content_version = _get_hub_model_version_for_marketplace_version(
-        hub_content_summaries, hub_model_version
-    )
-
     try:
         return _get_hub_model_version_for_open_weight_version(
             hub_content_summaries, hub_model_version
         )
     except KeyError as e:
+        marketplace_hub_content_version = _get_hub_model_version_for_marketplace_version(
+            hub_content_summaries, hub_model_version
+        )
         if marketplace_hub_content_version:
             return marketplace_hub_content_version
-        raise e
+        raise
 
 
 def _get_hub_model_version_for_open_weight_version(

tests/unit/sagemaker/jumpstart/constants.py

@@ -9178,6 +9178,7 @@
         "TrainingArtifactS3DataType": "S3Prefix",
         "TrainingArtifactCompressionType": "None",
         "TrainingArtifactUri": "s3://jumpstart-cache-prod-us-west-2/huggingface-training/train-huggingface-llm-gemma-2b-instruct.tar.gz",  # noqa: E501
+        "ModelTypes": ["OPEN_WEIGHTS", "PROPRIETARY"],
         "Hyperparameters": [
             {
                 "Name": "peft_type",

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

@@ -13,6 +13,10 @@
 from __future__ import absolute_import
 import pytest
 from sagemaker.jumpstart.hub.parser_utils import camel_to_snake
+from sagemaker.jumpstart.hub.parsers import make_model_specs_from_describe_hub_content_response
+from sagemaker.jumpstart.hub.interfaces import HubModelDocument
+from tests.unit.sagemaker.jumpstart.constants import HUB_MODEL_DOCUMENT_DICTS
+from unittest.mock import MagicMock
 
 REGION = "us-east-1"
 ACCOUNT_ID = "123456789123"
@@ -32,3 +36,14 @@
 )
 def test_parse_(input_string, expected):
     assert expected == camel_to_snake(input_string)
+
+
+def test_make_model_specs_from_describe_hub_content_response():
+    mock_describe_response = MagicMock()
+    region = "us-west-2"
+    mock_describe_response.get_hub_region.return_value = region
+    mock_describe_response.hub_content_version = "1.0.0"
+    json_obj = HUB_MODEL_DOCUMENT_DICTS["huggingface-llm-gemma-2b-instruct"]
+    mock_describe_response.hub_content_document = HubModelDocument(json_obj=json_obj, region=region)
+
+    make_model_specs_from_describe_hub_content_response(mock_describe_response)