@@ -1722,10 +1722,8 @@ def test_build_mlflow_model_local_input_happy(
1722
1722
mock_path_exists .side_effect = lambda path : True if path == "test_path" else False
1723
1723
1724
1724
mock_mode = Mock ()
1725
- mock_sageMakerEndpointMode .side_effect = (
1726
- lambda inference_spec , model_server : mock_mode
1727
- if inference_spec is None and model_server == ModelServer .TORCHSERVE
1728
- else None
1725
+ mock_sageMakerEndpointMode .side_effect = lambda inference_spec , model_server : (
1726
+ mock_mode if inference_spec is None and model_server == ModelServer .TORCHSERVE else None
1729
1727
)
1730
1728
mock_mode .prepare .return_value = (
1731
1729
model_data ,
@@ -1799,10 +1797,8 @@ def test_build_mlflow_model_s3_input_happy(
1799
1797
mock_path_exists .side_effect = lambda path : True if path == "test_path" else False
1800
1798
1801
1799
mock_mode = Mock ()
1802
- mock_sageMakerEndpointMode .side_effect = (
1803
- lambda inference_spec , model_server : mock_mode
1804
- if inference_spec is None and model_server == ModelServer .TORCHSERVE
1805
- else None
1800
+ mock_sageMakerEndpointMode .side_effect = lambda inference_spec , model_server : (
1801
+ mock_mode if inference_spec is None and model_server == ModelServer .TORCHSERVE else None
1806
1802
)
1807
1803
mock_mode .prepare .return_value = (
1808
1804
model_data ,
@@ -1875,10 +1871,8 @@ def test_build_mlflow_model_s3_input_non_mlflow_case(
1875
1871
mock_path_exists .side_effect = lambda path : True if path == "test_path" else False
1876
1872
1877
1873
mock_mode = Mock ()
1878
- mock_sageMakerEndpointMode .side_effect = (
1879
- lambda inference_spec , model_server : mock_mode
1880
- if inference_spec is None and model_server == ModelServer .TORCHSERVE
1881
- else None
1874
+ mock_sageMakerEndpointMode .side_effect = lambda inference_spec , model_server : (
1875
+ mock_mode if inference_spec is None and model_server == ModelServer .TORCHSERVE else None
1882
1876
)
1883
1877
mock_mode .prepare .return_value = (
1884
1878
model_data ,
0 commit comments