only change some called_with into assert_called_with according to py312 unit-tests

Author: mollyheamazon

tests/unit/sagemaker/modules/train/test_model_trainer.py

@@ -1048,7 +1048,7 @@ def mock_upload_data(path, bucket, key_prefix):
 
     model_trainer.train()
 
-    assert mock_local_container.train.called_once_with(
+    mock_local_container.train.assert_called_once_with(
         training_job_name=unique_name,
         instance_type=compute.instance_type,
         instance_count=compute.instance_count,

tests/unit/sagemaker/workflow/test_pipeline.py

@@ -99,7 +99,7 @@ def test_pipeline_create_and_update_with_config_injection(sagemaker_session_mock
         RoleArn=pipeline_role_arn,
     )
     pipeline.upsert()
-    assert sagemaker_session_mock.sagemaker_client.update_pipeline.called_with(
+    sagemaker_session_mock.sagemaker_client.update_pipeline.assert_called_with(
         PipelineName="MyPipeline",
         PipelineDefinition=pipeline.definition(),
         RoleArn=pipeline_role_arn,
@@ -130,7 +130,7 @@ def test_pipeline_create_with_parallelism_config(sagemaker_session_mock, role_ar
         role_arn=role_arn,
         parallelism_config=dict(MaxParallelExecutionSteps=10),
     )
-    assert sagemaker_session_mock.sagemaker_client.create_pipeline.called_with(
+    sagemaker_session_mock.sagemaker_client.create_pipeline.assert_called_with(
         PipelineName="MyPipeline",
         PipelineDefinition=pipeline.definition(),
         RoleArn=role_arn,
@@ -149,7 +149,7 @@ def test_pipeline_create_and_start_with_parallelism_config(sagemaker_session_moc
         role_arn=role_arn,
         parallelism_config=dict(MaxParallelExecutionSteps=10),
     )
-    assert sagemaker_session_mock.sagemaker_client.create_pipeline.called_with(
+    sagemaker_session_mock.sagemaker_client.create_pipeline.assert_called_with(
         PipelineName="MyPipeline",
         PipelineDefinition=pipeline.definition(),
         RoleArn=role_arn,
@@ -209,7 +209,7 @@ def test_pipeline_update(sagemaker_session_mock, role_arn):
     assert not pipeline.steps
     pipeline.update(role_arn=role_arn)
     assert len(json.loads(pipeline.definition())["Steps"]) == 0
-    assert sagemaker_session_mock.sagemaker_client.update_pipeline.called_with(
+    sagemaker_session_mock.sagemaker_client.update_pipeline.assert_called_with(
         PipelineName="MyPipeline", PipelineDefinition=pipeline.definition(), RoleArn=role_arn
     )
 
@@ -345,7 +345,7 @@ def test_pipeline_update_with_parallelism_config(sagemaker_session_mock, role_ar
         role_arn=role_arn,
         parallelism_config=dict(MaxParallelExecutionSteps=10),
     )
-    assert sagemaker_session_mock.sagemaker_client.update_pipeline.called_with(
+    sagemaker_session_mock.sagemaker_client.update_pipeline.assert_called_with(
         PipelineName="MyPipeline",
         PipelineDefinition=pipeline.definition(),
         RoleArn=role_arn,
@@ -418,7 +418,7 @@ def _raise_does_already_exists_client_error(**kwargs):
     sagemaker_session_mock.sagemaker_client.update_pipeline.assert_called_once_with(
         PipelineName="MyPipeline", PipelineDefinition=pipeline.definition(), RoleArn=role_arn
     )
-    assert sagemaker_session_mock.sagemaker_client.list_tags.called_with(
+    sagemaker_session_mock.sagemaker_client.list_tags.assert_called_with(
         ResourceArn="mock_pipeline_arn"
     )
 
@@ -523,7 +523,7 @@ def test_pipeline_delete(sagemaker_session_mock):
         sagemaker_session=sagemaker_session_mock,
     )
     pipeline.delete()
-    assert sagemaker_session_mock.sagemaker_client.delete_pipeline.called_with(
+    sagemaker_session_mock.sagemaker_client.delete_pipeline.assert_called_with(
         PipelineName="MyPipeline",
     )
 
@@ -536,7 +536,7 @@ def test_pipeline_describe(sagemaker_session_mock):
         sagemaker_session=sagemaker_session_mock,
     )
     pipeline.describe()
-    assert sagemaker_session_mock.sagemaker_client.describe_pipeline.called_with(
+    sagemaker_session_mock.sagemaker_client.describe_pipeline.assert_called_with(
         PipelineName="MyPipeline",
     )
 
@@ -552,7 +552,7 @@ def test_pipeline_start(sagemaker_session_mock):
         sagemaker_session=sagemaker_session_mock,
     )
     pipeline.start()
-    assert sagemaker_session_mock.start_pipeline_execution.called_with(
+    sagemaker_session_mock.start_pipeline_execution.assert_called_with(
         PipelineName="MyPipeline",
     )
 
@@ -821,10 +821,8 @@ def test_pipeline_build_parameters_from_execution(sagemaker_session_mock):
         pipeline_execution_arn=reference_execution_arn,
         parameter_value_overrides=parameter_value_overrides,
     )
-    assert (
-        sagemaker_session_mock.sagemaker_client.list_pipeline_parameters_for_execution.called_with(
-            PipelineExecutionArn=reference_execution_arn
-        )
+    sagemaker_session_mock.sagemaker_client.list_pipeline_parameters_for_execution.assert_called_with(
+        PipelineExecutionArn=reference_execution_arn
     )
     assert len(parameters) == 1
     assert parameters["TestParameterName"] == "NewParameterValue"
@@ -850,10 +848,8 @@ def test_pipeline_build_parameters_from_execution_with_invalid_overrides(sagemak
         + f"are not present in the pipeline execution: {reference_execution_arn}"
         in str(error)
     )
-    assert (
-        sagemaker_session_mock.sagemaker_client.list_pipeline_parameters_for_execution.called_with(
-            PipelineExecutionArn=reference_execution_arn
-        )
+    sagemaker_session_mock.sagemaker_client.list_pipeline_parameters_for_execution.assert_called_with(
+        PipelineExecutionArn=reference_execution_arn
     )
 
 
@@ -912,7 +908,7 @@ def test_pipeline_execution_basics(sagemaker_session_mock):
         PipelineExecutionArn="my:arn"
     )
     execution.describe()
-    assert sagemaker_session_mock.sagemaker_client.describe_pipeline_execution.called_with(
+    sagemaker_session_mock.sagemaker_client.describe_pipeline_execution.assert_called_with(
         PipelineExecutionArn="my:arn"
     )
     steps = execution.list_steps()