fix black

Author: mollyheamazon

tests/integ/sagemaker/serve/test_serve_js_deep_unit_tests.py

@@ -24,11 +24,12 @@
 def test_js_model_with_optimize_speculative_decoding_config_gated_requests_are_expected(
     sagemaker_session,
 ):
-    with patch.object(
-        Session, "create_model", return_value="mock_model"
-    ) as mock_create_model, patch.object(
-        Session, "endpoint_from_production_variants"
-    ) as mock_endpoint_from_production_variants:
+    with (
+        patch.object(Session, "create_model", return_value="mock_model") as mock_create_model,
+        patch.object(
+            Session, "endpoint_from_production_variants"
+        ) as mock_endpoint_from_production_variants,
+    ):
         iam_client = sagemaker_session.boto_session.client("iam")
         role_arn = iam_client.get_role(RoleName=ROLE_NAME)["Role"]["Arn"]
 
@@ -100,17 +101,18 @@ def test_js_model_with_optimize_speculative_decoding_config_gated_requests_are_e
 def test_js_model_with_optimize_sharding_and_resource_requirements_requests_are_expected(
     sagemaker_session,
 ):
-    with patch.object(
-        Session,
-        "wait_for_optimization_job",
-        return_value={"OptimizationJobName": "mock_optimization_job"},
-    ), patch.object(
-        Session, "create_model", return_value="mock_model"
-    ) as mock_create_model, patch.object(
-        Session, "endpoint_from_production_variants", return_value="mock_endpoint_name"
-    ) as mock_endpoint_from_production_variants, patch.object(
-        Session, "create_inference_component"
-    ) as mock_create_inference_component:
+    with (
+        patch.object(
+            Session,
+            "wait_for_optimization_job",
+            return_value={"OptimizationJobName": "mock_optimization_job"},
+        ),
+        patch.object(Session, "create_model", return_value="mock_model") as mock_create_model,
+        patch.object(
+            Session, "endpoint_from_production_variants", return_value="mock_endpoint_name"
+        ) as mock_endpoint_from_production_variants,
+        patch.object(Session, "create_inference_component") as mock_create_inference_component,
+    ):
         iam_client = sagemaker_session.boto_session.client("iam")
         role_arn = iam_client.get_role(RoleName=ROLE_NAME)["Role"]["Arn"]
 
@@ -185,15 +187,17 @@ def test_js_model_with_optimize_sharding_and_resource_requirements_requests_are_
 def test_js_model_with_optimize_quantization_on_pre_optimized_model_requests_are_expected(
     sagemaker_session,
 ):
-    with patch.object(
-        Session,
-        "wait_for_optimization_job",
-        return_value={"OptimizationJobName": "mock_optimization_job"},
-    ), patch.object(
-        Session, "create_model", return_value="mock_model"
-    ) as mock_create_model, patch.object(
-        Session, "endpoint_from_production_variants", return_value="mock_endpoint_name"
-    ) as mock_endpoint_from_production_variants:
+    with (
+        patch.object(
+            Session,
+            "wait_for_optimization_job",
+            return_value={"OptimizationJobName": "mock_optimization_job"},
+        ),
+        patch.object(Session, "create_model", return_value="mock_model") as mock_create_model,
+        patch.object(
+            Session, "endpoint_from_production_variants", return_value="mock_endpoint_name"
+        ) as mock_endpoint_from_production_variants,
+    ):
         iam_client = sagemaker_session.boto_session.client("iam")
         role_arn = iam_client.get_role(RoleName=ROLE_NAME)["Role"]["Arn"]