add important unit test

Author: bencrabtree

src/sagemaker/jumpstart/accessors.py

@@ -268,7 +268,7 @@ def get_model_specs(
             hub_model_arn = utils.construct_hub_model_arn_from_inputs(
                 hub_arn=hub_arn, model_name=model_id, version=version
             )
-            return JumpStartModelsAccessor._cache.get_hub_model(hub_model_arn)
+            return JumpStartModelsAccessor._cache.get_hub_model(hub_model_arn=hub_model_arn)
 
         return JumpStartModelsAccessor._cache.get_specs(  # type: ignore
             model_id=model_id, semantic_version_str=version

tests/unit/sagemaker/jumpstart/test_accessors.py

@@ -72,6 +72,35 @@ def test_jumpstart_models_cache_get_fxs(mock_cache):
     reload(accessors)
 
 
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor._cache")
+def test_jumpstart_models_cache_get_model_specs(mock_cache):
+    mock_cache.get_specs = Mock()
+    mock_cache.get_hub_model = Mock()
+    model_id, version = "pytorch-ic-mobilenet-v2", "*"
+    region = "us-west-2"
+
+    accessors.JumpStartModelsAccessor.get_model_specs(
+        region=region, model_id=model_id, version=version
+    )
+    mock_cache.get_specs.assert_called_once_with(model_id=model_id, semantic_version_str=version)
+    mock_cache.get_hub_model.assert_not_called()
+
+    accessors.JumpStartModelsAccessor.get_model_specs(
+        region=region,
+        model_id=model_id,
+        version=version,
+        hub_arn=f"arn:aws:sagemaker:{region}:123456789123:hub/my-mock-hub",
+    )
+    mock_cache.get_hub_model.assert_called_once_with(
+        hub_model_arn=(
+            f"arn:aws:sagemaker:{region}:123456789123:hub-content/my-mock-hub/Model/{model_id}/{version}"
+        )
+    )
+
+    # necessary because accessors is a static module
+    reload(accessors)
+
+
 @patch("sagemaker.jumpstart.cache.JumpStartModelsCache")
 def test_jumpstart_models_cache_set_reset_fxs(mock_model_cache: Mock):
 

tests/unit/sagemaker/jumpstart/utils.py

@@ -148,6 +148,7 @@ def get_spec_from_base_spec(
     semantic_version_str: str = None,
     version: str = None,
     hub_arn: Optional[str] = None,
+    hub_model_arn: Optional[str] = None,
     s3_client: boto3.client = None,
 ) -> JumpStartModelSpecs: