Fix unit tests

Author: Jonathan Makunga

tests/unit/test_session.py

@@ -3192,7 +3192,9 @@ def test_create_model_from_job_with_vpc_config_override(sagemaker_session):
 
 def test_endpoint_from_production_variants(sagemaker_session):
     ims = sagemaker_session
-    ims.sagemaker_client.describe_endpoint = Mock(return_value={"EndpointStatus": "InService"})
+    ims.sagemaker_client.describe_endpoint = Mock(
+        return_value={"EndpointStatus": "InService", "EndpointArn": "arn:aws:sagemaker:"}
+    )
     pvs = [
         sagemaker.production_variant("A", "ml.p2.xlarge"),
         sagemaker.production_variant("B", "p299.4096xlarge"),
@@ -3487,7 +3489,7 @@ def test_endpoint_from_production_variants_with_sagemaker_config_injection(
     sagemaker_session.sagemaker_config = SAGEMAKER_CONFIG_ENDPOINT_CONFIG
 
     sagemaker_session.sagemaker_client.describe_endpoint = Mock(
-        return_value={"EndpointStatus": "InService"}
+        return_value={"EndpointStatus": "InService", "EndpointArn": "arn:aws:sagemaker:"}
     )
     pvs = [
         sagemaker.production_variant("A", "ml.p2.xlarge"),
@@ -3555,7 +3557,7 @@ def test_endpoint_from_production_variants_with_sagemaker_config_injection_parti
     sagemaker_session.sagemaker_config = SAGEMAKER_CONFIG_ENDPOINT_CONFIG
 
     sagemaker_session.sagemaker_client.describe_endpoint = Mock(
-        return_value={"EndpointStatus": "InService"}
+        return_value={"EndpointStatus": "InService", "EndpointArn": "arn:aws:sagemaker:"}
     )
     pvs = [
         sagemaker.production_variant("A", "ml.g5.xlarge"),
@@ -3619,7 +3621,7 @@ def test_endpoint_from_production_variants_with_sagemaker_config_injection_no_km
     sagemaker_session.sagemaker_config = SAGEMAKER_CONFIG_ENDPOINT_CONFIG
 
     sagemaker_session.sagemaker_client.describe_endpoint = Mock(
-        return_value={"EndpointStatus": "InService"}
+        return_value={"EndpointStatus": "InService", "EndpointArn": "arn:aws:sagemaker:"}
     )
     pvs = [
         sagemaker.production_variant("A", "ml.g5.xlarge"),
@@ -3726,7 +3728,9 @@ def test_create_endpoint_config_with_explainer_config(sagemaker_session):
 
 def test_endpoint_from_production_variants_with_tags(sagemaker_session):
     ims = sagemaker_session
-    ims.sagemaker_client.describe_endpoint = Mock(return_value={"EndpointStatus": "InService"})
+    ims.sagemaker_client.describe_endpoint = Mock(
+        return_value={"EndpointStatus": "InService", "EndpointArn": "arn:aws:sagemaker:"}
+    )
     pvs = [
         sagemaker.production_variant("A", "ml.p2.xlarge"),
         sagemaker.production_variant("B", "p299.4096xlarge"),
@@ -3757,7 +3761,9 @@ def test_endpoint_from_production_variants_with_combined_sagemaker_config_inject
     sagemaker_session.sagemaker_config = SAGEMAKER_CONFIG_ENDPOINT_ENDPOINT_CONFIG_COMBINED
 
     ims = sagemaker_session
-    ims.sagemaker_client.describe_endpoint = Mock(return_value={"EndpointStatus": "InService"})
+    ims.sagemaker_client.describe_endpoint = Mock(
+        return_value={"EndpointStatus": "InService", "EndpointArn": "arn:aws:sagemaker:"}
+    )
     pvs = [
         sagemaker.production_variant("A", "ml.p2.xlarge"),
         sagemaker.production_variant("B", "p299.4096xlarge"),
@@ -3801,7 +3807,9 @@ def test_endpoint_from_production_variants_with_sagemaker_config_injection_tags(
     sagemaker_session.sagemaker_config = SAGEMAKER_CONFIG_ENDPOINT
 
     ims = sagemaker_session
-    ims.sagemaker_client.describe_endpoint = Mock(return_value={"EndpointStatus": "InService"})
+    ims.sagemaker_client.describe_endpoint = Mock(
+        return_value={"EndpointStatus": "InService", "EndpointArn": "arn:aws:sagemaker:"}
+    )
     pvs = [
         sagemaker.production_variant("A", "ml.p2.xlarge"),
         sagemaker.production_variant("B", "p299.4096xlarge"),
@@ -3830,7 +3838,9 @@ def test_endpoint_from_production_variants_with_sagemaker_config_injection_tags(
 
 def test_endpoint_from_production_variants_with_accelerator_type(sagemaker_session):
     ims = sagemaker_session
-    ims.sagemaker_client.describe_endpoint = Mock(return_value={"EndpointStatus": "InService"})
+    ims.sagemaker_client.describe_endpoint = Mock(
+        return_value={"EndpointStatus": "InService", "EndpointArn": "arn:aws:sagemaker:"}
+    )
     pvs = [
         sagemaker.production_variant("A", "ml.p2.xlarge", accelerator_type=ACCELERATOR_TYPE),
         sagemaker.production_variant("B", "p299.4096xlarge", accelerator_type=ACCELERATOR_TYPE),
@@ -3861,7 +3871,9 @@ def test_endpoint_from_production_variants_with_accelerator_type_sagemaker_confi
     sagemaker_session.sagemaker_config = SAGEMAKER_CONFIG_ENDPOINT
 
     ims = sagemaker_session
-    ims.sagemaker_client.describe_endpoint = Mock(return_value={"EndpointStatus": "InService"})
+    ims.sagemaker_client.describe_endpoint = Mock(
+        return_value={"EndpointStatus": "InService", "EndpointArn": "arn:aws:sagemaker:"}
+    )
     pvs = [
         sagemaker.production_variant("A", "ml.p2.xlarge", accelerator_type=ACCELERATOR_TYPE),
         sagemaker.production_variant("B", "p299.4096xlarge", accelerator_type=ACCELERATOR_TYPE),
@@ -3892,7 +3904,9 @@ def test_endpoint_from_production_variants_with_serverless_inference_config(
     sagemaker_session,
 ):
     ims = sagemaker_session
-    ims.sagemaker_client.describe_endpoint = Mock(return_value={"EndpointStatus": "InService"})
+    ims.sagemaker_client.describe_endpoint = Mock(
+        return_value={"EndpointStatus": "InService", "EndpointArn": "arn:aws:sagemaker:"}
+    )
     pvs = [
         sagemaker.production_variant(
             "A", "ml.p2.xlarge", serverless_inference_config=SERVERLESS_INFERENCE_CONFIG
@@ -3929,7 +3943,9 @@ def test_endpoint_from_production_variants_with_serverless_inference_config_sage
     sagemaker_session.sagemaker_config = SAGEMAKER_CONFIG_ENDPOINT
 
     ims = sagemaker_session
-    ims.sagemaker_client.describe_endpoint = Mock(return_value={"EndpointStatus": "InService"})
+    ims.sagemaker_client.describe_endpoint = Mock(
+        return_value={"EndpointStatus": "InService", "EndpointArn": "arn:aws:sagemaker:"}
+    )
     pvs = [
         sagemaker.production_variant(
             "A", "ml.p2.xlarge", serverless_inference_config=SERVERLESS_INFERENCE_CONFIG
@@ -3964,7 +3980,9 @@ def test_endpoint_from_production_variants_with_serverless_inference_config_sage
 
 def test_endpoint_from_production_variants_with_async_config(sagemaker_session):
     ims = sagemaker_session
-    ims.sagemaker_client.describe_endpoint = Mock(return_value={"EndpointStatus": "InService"})
+    ims.sagemaker_client.describe_endpoint = Mock(
+        return_value={"EndpointStatus": "InService", "EndpointArn": "arn:aws:sagemaker:"}
+    )
     pvs = [
         sagemaker.production_variant("A", "ml.p2.xlarge"),
         sagemaker.production_variant("B", "p299.4096xlarge"),
@@ -4000,7 +4018,9 @@ def test_endpoint_from_production_variants_with_async_config_sagemaker_config_in
     sagemaker_session.sagemaker_config = SAGEMAKER_CONFIG_ENDPOINT
 
     ims = sagemaker_session
-    ims.sagemaker_client.describe_endpoint = Mock(return_value={"EndpointStatus": "InService"})
+    ims.sagemaker_client.describe_endpoint = Mock(
+        return_value={"EndpointStatus": "InService", "EndpointArn": "arn:aws:sagemaker:"}
+    )
     pvs = [
         sagemaker.production_variant("A", "ml.p2.xlarge"),
         sagemaker.production_variant("B", "p299.4096xlarge"),
@@ -4037,7 +4057,9 @@ def test_endpoint_from_production_variants_with_clarify_explainer_config(
     sagemaker_session,
 ):
     ims = sagemaker_session
-    ims.sagemaker_client.describe_endpoint = Mock(return_value={"EndpointStatus": "InService"})
+    ims.sagemaker_client.describe_endpoint = Mock(
+        return_value={"EndpointStatus": "InService", "EndpointArn": "arn:aws:sagemaker:"}
+    )
     pvs = [
         sagemaker.production_variant("A", "ml.p2.xlarge"),
         sagemaker.production_variant("B", "p299.4096xlarge"),
@@ -4069,7 +4091,7 @@ def test_endpoint_from_production_variants_with_clarify_explainer_config(
 
 def test_update_endpoint_succeed(sagemaker_session):
     sagemaker_session.sagemaker_client.describe_endpoint = Mock(
-        return_value={"EndpointStatus": "InService"}
+        return_value={"EndpointStatus": "InService", "EndpointArn": "arn:aws:sagemaker:"}
     )
     endpoint_name = "some-endpoint"
     endpoint_config = "some-endpoint-config"
@@ -4079,7 +4101,7 @@ def test_update_endpoint_succeed(sagemaker_session):
 
 def test_update_endpoint_no_wait(sagemaker_session):
     sagemaker_session.sagemaker_client.describe_endpoint = Mock(
-        return_value={"EndpointStatus": "Updating"}
+        return_value={"EndpointStatus": "Updating", "EndpointArn": "arn:aws:sagemaker:"}
     )
     endpoint_name = "some-endpoint"
     endpoint_config = "some-endpoint-config"
@@ -6136,7 +6158,10 @@ def test_upload_data_default_bucket_and_prefix_combinations(
 
 def test_is_inference_component_based_endpoint_affirmative(sagemaker_session):
 
-    describe_endpoint_response = {"EndpointConfigName": "some-endpoint-config"}
+    describe_endpoint_response = {
+        "EndpointConfigName": "some-endpoint-config",
+        "EndpointArn": "arn:aws:sagemaker:",
+    }
     describe_endpoint_config_response = {
         "ExecutionRoleArn": "some-role-arn",
         "ProductionVariants": [{"VariantName": "AllTraffic"}],
@@ -6160,7 +6185,10 @@ def test_is_inference_component_based_endpoint_affirmative(sagemaker_session):
 
 def test_is_inference_component_based_endpoint_negative_no_role(sagemaker_session):
 
-    describe_endpoint_response = {"EndpointConfigName": "some-endpoint-config"}
+    describe_endpoint_response = {
+        "EndpointConfigName": "some-endpoint-config",
+        "EndpointArn": "arn:aws:sagemaker:",
+    }
     describe_endpoint_config_response = {
         "ProductionVariants": [{"VariantName": "AllTraffic"}],
     }
@@ -6183,7 +6211,10 @@ def test_is_inference_component_based_endpoint_negative_no_role(sagemaker_sessio
 
 def test_is_inference_component_based_endpoint_positive_multiple_variants(sagemaker_session):
 
-    describe_endpoint_response = {"EndpointConfigName": "some-endpoint-config"}
+    describe_endpoint_response = {
+        "EndpointConfigName": "some-endpoint-config",
+        "EndpointArn": "arn:aws:sagemaker:",
+    }
     describe_endpoint_config_response = {
         "ExecutionRoleArn": "some-role-arn",
         "ProductionVariants": [{"VariantName": "AllTraffic1"}, {"VariantName": "AllTraffic2"}],
@@ -6207,7 +6238,10 @@ def test_is_inference_component_based_endpoint_positive_multiple_variants(sagema
 
 def test_is_inference_component_based_endpoint_negative_no_variants(sagemaker_session):
 
-    describe_endpoint_response = {"EndpointConfigName": "some-endpoint-config"}
+    describe_endpoint_response = {
+        "EndpointConfigName": "some-endpoint-config",
+        "EndpointArn": "arn:aws:sagemaker:",
+    }
     describe_endpoint_config_response = {
         "ExecutionRoleArn": "some-role-arn",
         "ProductionVariants": [],
@@ -6231,7 +6265,10 @@ def test_is_inference_component_based_endpoint_negative_no_variants(sagemaker_se
 
 def test_is_inference_component_based_endpoint_negative_model_name_present(sagemaker_session):
 
-    describe_endpoint_response = {"EndpointConfigName": "some-endpoint-config"}
+    describe_endpoint_response = {
+        "EndpointConfigName": "some-endpoint-config",
+        "EndpointArn": "arn:aws:sagemaker:",
+    }
     describe_endpoint_config_response = {
         "ExecutionRoleArn": "some-role-arn",
         "ProductionVariants": [