fix: Adding unit tests

Author: chrstfu

src/sagemaker/jumpstart/curated_hub/types.py

@@ -27,7 +27,7 @@ class CuratedHubTagName(str, Enum):
 
 @dataclass
 class Tag:
-  key: str
+  key: CuratedHubTagName
   value: str
 
 @dataclass

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -196,15 +196,18 @@ def tag_hub_content(hub_content_arn: str, tags: List[Tag], session: Session) ->
     return responses
 
 def find_all_tags_for_jumpstart_model(hub_name: str, hub_content_name: str, region: str, session: Session) -> List[Tag]:
-    hub_content_versions = session.list_hub_content_versions(
+    list_versions_response = session.list_hub_content_versions(
         hub_name=hub_name,
         hub_content_type='Model',
         hub_content_name=hub_content_name
     )
+    hub_content_versions = list_versions_response["HubContentSummaries"]
 
     tag_name_to_versions_map: Dict[CuratedHubTagName, List[str]] = {}
     for hub_content_version_summary in hub_content_versions:
         jumpstart_model = get_jumpstart_model_and_version(hub_content_version_summary)
+        if jumpstart_model["model_id"] is None or jumpstart_model["version"] is None:
+            continue
         tag_names_to_add: List[CuratedHubTagName] = find_tags_for_jumpstart_model_version(
             model_id=jumpstart_model["model_id"],
             version=jumpstart_model["version"],
@@ -215,13 +218,13 @@ def find_all_tags_for_jumpstart_model(hub_name: str, hub_content_name: str, regi
         for tag_name in tag_names_to_add:
           if tag_name not in tag_name_to_versions_map:
               tag_name_to_versions_map[tag_name] = []
-          tag_name_to_versions_map[tag_name].append(jumpstart_model["version"])
+          tag_name_to_versions_map[tag_name].append(hub_content_version_summary["HubContentVersion"])
 
     tags: List[Tag] = []
-    for tag_name, versions in tag_name_to_versions_map:
+    for tag_name, versions in tag_name_to_versions_map.items():
         tags.append(Tag(
             key=tag_name,
-            versions=str(versions)
+            value=str(versions)
         ))
 
     return tags
@@ -241,11 +244,11 @@ def find_tags_for_jumpstart_model_version(model_id: str, version: str, region: s
     )
 
     if (specs.deprecated):
-          tags_to_add.add(CuratedHubTagName.DEPRECATED_VERSIONS_TAG)
+          tags_to_add.append(CuratedHubTagName.DEPRECATED_VERSIONS_TAG)
     if (specs.inference_vulnerable):
-        tags_to_add.add(CuratedHubTagName.INFERENCE_VULNERABLE_VERSIONS_TAG)
+        tags_to_add.append(CuratedHubTagName.INFERENCE_VULNERABLE_VERSIONS_TAG)
     if (specs.training_vulnerable):
-        tags_to_add.add(CuratedHubTagName.TRAINING_VULNERABLE_VERSIONS_TAG)
+        tags_to_add.append(CuratedHubTagName.TRAINING_VULNERABLE_VERSIONS_TAG)
 
     return tags_to_add
 

tests/unit/sagemaker/jumpstart/curated_hub/test_utils.py

@@ -15,7 +15,13 @@
 from unittest.mock import Mock
 from sagemaker.jumpstart.types import HubArnExtractedInfo
 from sagemaker.jumpstart.constants import JUMPSTART_DEFAULT_REGION_NAME
+from sagemaker.jumpstart.enums import JumpStartScriptScope
 from sagemaker.jumpstart.curated_hub import utils
+from unittest.mock import patch
+from sagemaker.jumpstart.curated_hub.types import (
+    Tag,
+    CuratedHubTagName
+)
 
 
 def test_get_info_from_hub_resource_arn():
@@ -168,3 +174,110 @@ def test_create_hub_bucket_if_it_does_not_exist():
 
     mock_sagemaker_session.boto_session.resource("s3").create_bucketassert_called_once()
     assert created_hub_bucket_name == bucket_name
+
+@patch("sagemaker.jumpstart.utils.verify_model_region_and_return_specs")
+def test_find_tags_for_jumpstart_model_version(mock_spec_util):
+    mock_sagemaker_session = Mock()
+    mock_specs = Mock()
+    mock_specs.deprecated = True
+    mock_specs.inference_vulnerable = True
+    mock_specs.training_vulnerable = True
+    mock_spec_util.return_value = mock_specs
+
+    tags = utils.find_tags_for_jumpstart_model_version(
+        model_id="test",
+        version="test",
+        region="test",
+        session=mock_sagemaker_session
+    )
+
+    mock_spec_util.assert_called_once_with(
+        model_id="test",
+        version="test",
+        region="test",
+        scope=JumpStartScriptScope.INFERENCE,
+        tolerate_vulnerable_model = True,
+        tolerate_deprecated_model = True,
+        sagemaker_session=mock_sagemaker_session,
+    )
+
+    assert tags == [CuratedHubTagName.DEPRECATED_VERSIONS_TAG, CuratedHubTagName.INFERENCE_VULNERABLE_VERSIONS_TAG, CuratedHubTagName.TRAINING_VULNERABLE_VERSIONS_TAG]
+
+@patch("sagemaker.jumpstart.utils.verify_model_region_and_return_specs")
+def test_find_tags_for_jumpstart_model_version_some_false(mock_spec_util):
+    mock_sagemaker_session = Mock()
+    mock_specs = Mock()
+    mock_specs.deprecated = True
+    mock_specs.inference_vulnerable = False
+    mock_specs.training_vulnerable = False
+    mock_spec_util.return_value = mock_specs
+
+    tags = utils.find_tags_for_jumpstart_model_version(
+        model_id="test",
+        version="test",
+        region="test",
+        session=mock_sagemaker_session
+    )
+
+    mock_spec_util.assert_called_once_with(
+        model_id="test",
+        version="test",
+        region="test",
+        scope=JumpStartScriptScope.INFERENCE,
+        tolerate_vulnerable_model = True,
+        tolerate_deprecated_model = True,
+        sagemaker_session=mock_sagemaker_session,
+    )
+
+    assert tags == [CuratedHubTagName.DEPRECATED_VERSIONS_TAG]
+
+@patch("sagemaker.jumpstart.utils.verify_model_region_and_return_specs")
+def test_find_all_tags_for_jumpstart_model(mock_spec_util):
+    mock_sagemaker_session = Mock()
+    mock_sagemaker_session.list_hub_content_versions.return_value = {
+        "HubContentSummaries": [
+            {
+                "HubContentVersion": "1.0.0",
+                "search_keywords": [
+                    "@jumpstart-model-id:model-one-pytorch",
+                    "@jumpstart-model-version:1.0.3",
+                ]
+            },
+            {
+                "HubContentVersion": "2.0.0",
+                "search_keywords": [
+                  "@jumpstart-model-id:model-four-huggingface",
+                  "@jumpstart-model-version:2.0.2",
+                ]
+            },
+            {
+                "HubContentVersion": "3.0.0",
+                "search_keywords": []
+            }
+        ]
+    }
+
+    mock_specs = Mock()
+    mock_specs.deprecated = True
+    mock_specs.inference_vulnerable = True
+    mock_specs.training_vulnerable = True
+    mock_spec_util.return_value = mock_specs
+
+    tags = utils.find_all_tags_for_jumpstart_model(
+        hub_name="test",
+        hub_content_name="test",
+        region="test",
+        session=mock_sagemaker_session
+    )
+
+    mock_sagemaker_session.list_hub_content_versions.assert_called_once_with(
+        hub_name="test",
+        hub_content_type='Model',
+        hub_content_name="test",
+    )
+
+    assert tags == [
+        Tag(key=CuratedHubTagName.DEPRECATED_VERSIONS_TAG, value=str(["1.0.0", "2.0.0"])),
+        Tag(key=CuratedHubTagName.INFERENCE_VULNERABLE_VERSIONS_TAG, value=str(["1.0.0", "2.0.0"])),
+        Tag(key=CuratedHubTagName.TRAINING_VULNERABLE_VERSIONS_TAG, value=str(["1.0.0", "2.0.0"]))
+    ]