fix: tests

Author: chrstfu

src/sagemaker/jumpstart/curated_hub/utils.py

@@ -238,9 +238,9 @@ def _convert_to_tag_to_versions_map(
     for version, tags in version_to_tags_map.items():
         for tag in tags:
             if tag not in unsupported_hub_content_versions_map:
-                unsupported_hub_content_versions_map[tag] = []
+                unsupported_hub_content_versions_map[tag.value] = []
             # Versions for a HubContent are unique
-            unsupported_hub_content_versions_map[tag].append(version)
+            unsupported_hub_content_versions_map[tag.value].append(version)
 
     return unsupported_hub_content_versions_map
 

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

@@ -325,15 +325,15 @@ def test_find_all_tags_for_jumpstart_model_filters_non_jumpstart_models(mock_spe
 
     assert tags == [
         {
-            "Key": CuratedHubUnsupportedFlag.DEPRECATED_VERSIONS.value,
+            "Key": "deprecated_versions",
             "Value": str(["1.0.0", "2.0.0"]),
         },
         {
-            "Key": CuratedHubUnsupportedFlag.INFERENCE_VULNERABLE_VERSIONS.value,
+            "Key": "inference_vulnerable_versions",
             "Value": str(["1.0.0", "2.0.0"]),
         },
         {
-            "Key": CuratedHubUnsupportedFlag.TRAINING_VULNERABLE_VERSIONS.value,
+            "Key": "training_vulnerable_versions",
             "Value": str(["1.0.0", "2.0.0"]),
         },
     ]