chore: add additional unit test and fix formatting

Author: evakravi

tests/unit/sagemaker/jumpstart/test_utils.py

@@ -216,15 +216,33 @@ def test_add_jumpstart_tags_inference():
         inference_script_uri=inference_script_uri,
     ) == [{"Key": JumpStartTag.INFERENCE_MODEL_URI.value, "Value": inference_model_uri}]
 
-
     tags = []
     inference_model_uri = {"S3DataSource": {"S3Uri": random_jumpstart_s3_uri("random_key")}}
     inference_script_uri = "dfsdfs"
     assert utils.add_jumpstart_tags(
         tags=tags,
         inference_model_uri=inference_model_uri,
         inference_script_uri=inference_script_uri,
-    ) == [{"Key": JumpStartTag.INFERENCE_MODEL_URI.value, "Value": inference_model_uri["S3DataSource"]["S3Uri"]}]
+    ) == [
+        {
+            "Key": JumpStartTag.INFERENCE_MODEL_URI.value,
+            "Value": inference_model_uri["S3DataSource"]["S3Uri"],
+        }
+    ]
+
+    tags = []
+    inference_model_uri = {"S3DataSource": {"S3Uri": random_jumpstart_s3_uri("random_key/prefix/")}}
+    inference_script_uri = "dfsdfs"
+    assert utils.add_jumpstart_tags(
+        tags=tags,
+        inference_model_uri=inference_model_uri,
+        inference_script_uri=inference_script_uri,
+    ) == [
+        {
+            "Key": JumpStartTag.INFERENCE_MODEL_URI.value,
+            "Value": inference_model_uri["S3DataSource"]["S3Uri"],
+        }
+    ]
 
     tags = [{"Key": "some", "Value": "tag"}]
     inference_model_uri = random_jumpstart_s3_uri("random_key")