diff --git a/tests/unit/sagemaker/serve/builder/test_model_builder.py b/tests/unit/sagemaker/serve/builder/test_model_builder.py
index 67580744a9..85f178eb37 100644
--- a/tests/unit/sagemaker/serve/builder/test_model_builder.py
+++ b/tests/unit/sagemaker/serve/builder/test_model_builder.py
@@ -1722,10 +1722,8 @@ def test_build_mlflow_model_local_input_happy(
         mock_path_exists.side_effect = lambda path: True if path == "test_path" else False
 
         mock_mode = Mock()
-        mock_sageMakerEndpointMode.side_effect = (
-            lambda inference_spec, model_server: mock_mode
-            if inference_spec is None and model_server == ModelServer.TORCHSERVE
-            else None
+        mock_sageMakerEndpointMode.side_effect = lambda inference_spec, model_server: (
+	        mock_mode if inference_spec is None and model_server == ModelServer.TORCHSERVE else None
         )
         mock_mode.prepare.return_value = (
             model_data,
@@ -1799,10 +1797,8 @@ def test_build_mlflow_model_s3_input_happy(
         mock_path_exists.side_effect = lambda path: True if path == "test_path" else False
 
         mock_mode = Mock()
-        mock_sageMakerEndpointMode.side_effect = (
-            lambda inference_spec, model_server: mock_mode
-            if inference_spec is None and model_server == ModelServer.TORCHSERVE
-            else None
+        mock_sageMakerEndpointMode.side_effect = lambda inference_spec, model_server: (
+            mock_mode if inference_spec is None and model_server == ModelServer.TORCHSERVE else None
         )
         mock_mode.prepare.return_value = (
             model_data,
@@ -1875,10 +1871,8 @@ def test_build_mlflow_model_s3_input_non_mlflow_case(
         mock_path_exists.side_effect = lambda path: True if path == "test_path" else False
 
         mock_mode = Mock()
-        mock_sageMakerEndpointMode.side_effect = (
-            lambda inference_spec, model_server: mock_mode
-            if inference_spec is None and model_server == ModelServer.TORCHSERVE
-            else None
+        mock_sageMakerEndpointMode.side_effect = lambda inference_spec, model_server: (
+            mock_mode if inference_spec is None and model_server == ModelServer.TORCHSERVE else None
         )
         mock_mode.prepare.return_value = (
             model_data,