try changing some tests

Author: mollyheamazon

doc/requirements.txt

@@ -1,7 +1,7 @@
 sphinx==5.1.1
 sphinx-rtd-theme==0.5.0
 docutils==0.15.2
-packaging==20.9
+packaging>=23.0,<25
 jinja2==3.1.4
 schema==0.7.5
 accelerate>=0.24.1,<=0.27.0

pyproject.toml

@@ -41,7 +41,7 @@ dependencies = [
   "jsonschema",
   "numpy>=1.26.0,<2.0",
   "omegaconf>=2.2,<=2.3",
-  "packaging>=20.0",
+  "packaging>=23.0,<25",
   "pandas",
   "pathos",
   "platformdirs",

src/sagemaker/serve/utils/conda_in_process.yml

@@ -16,7 +16,7 @@ dependencies:
     - protobuf>=3.12,<5.0
     - smdebug_rulesconfig==1.0.1
     - importlib-metadata>=1.4.0,<7.0
-    - packaging>=20.0
+    - packaging>=23.0,<25
     - pandas
     - pathos
     - schema

tests/data/serve_resources/mlflow/pytorch/requirements.txt

@@ -7,7 +7,7 @@ dill==0.3.9
 gmpy2==2.1.2
 numpy==1.26.4
 opt-einsum==3.3.0
-packaging==21.3
+packaging>=23.0,<25
 pandas==2.2.1
 pyyaml==6.0.1
 requests==2.32.2

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,
@@ -168,7 +168,7 @@ def test_pipeline_create_and_start_with_parallelism_config(sagemaker_session_moc
 
     # Specify ParallelismConfiguration to another value which will be honored in backend
     pipeline.start(parallelism_config=dict(MaxParallelExecutionSteps=20))
-    assert sagemaker_session_mock.sagemaker_client.start_pipeline_execution.called_with(
+    sagemaker_session_mock.sagemaker_client.start_pipeline_execution.assert_called_with(
         PipelineName="MyPipeline",
         ParallelismConfiguration={"MaxParallelExecutionSteps": 20},
     )
@@ -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
     )
 
@@ -253,7 +253,7 @@ def test_pipeline_update(sagemaker_session_mock, role_arn):
 
     pipeline.update(role_arn=role_arn)
     assert len(json.loads(pipeline.definition())["Steps"]) == 3
-    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,13 +418,13 @@ 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(
-        ResourceArn="mock_pipeline_arn"
+    sagemaker_session_mock.sagemaker_client.list_tags.assert_called_with(
+        ResourceArn="pipeline-arn"
     )
 
     tags.append({"Key": "dummy", "Value": "dummy_tag"})
-    assert sagemaker_session_mock.sagemaker_client.add_tags.called_with(
-        ResourceArn="mock_pipeline_arn", Tags=tags
+    sagemaker_session_mock.sagemaker_client.add_tags.assert_called_with(
+        ResourceArn="pipeline-arn", Tags=tags
     )
 
 
@@ -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,17 +552,17 @@ 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",
     )
 
     pipeline.start(execution_display_name="pipeline-execution")
-    assert sagemaker_session_mock.start_pipeline_execution.called_with(
+    sagemaker_session_mock.start_pipeline_execution.assert_called_with(
         PipelineName="MyPipeline", PipelineExecutionDisplayName="pipeline-execution"
     )
 
     pipeline.start(parameters=dict(alpha="epsilon"))
-    assert sagemaker_session_mock.start_pipeline_execution.called_with(
+    sagemaker_session_mock.start_pipeline_execution.assert_called_with(
         PipelineName="MyPipeline", PipelineParameters=[{"Name": "alpha", "Value": "epsilon"}]
     )
 
@@ -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
     )
 
 
@@ -908,24 +904,23 @@ def test_pipeline_execution_basics(sagemaker_session_mock):
     )
     execution = pipeline.start()
     execution.stop()
-    assert sagemaker_session_mock.sagemaker_client.stop_pipeline_execution.called_with(
+    sagemaker_session_mock.sagemaker_client.stop_pipeline_execution.assert_called_with(
         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()
-    assert sagemaker_session_mock.sagemaker_client.describe_pipeline_execution_steps.called_with(
+    sagemaker_session_mock.sagemaker_client.describe_pipeline_execution_steps.assert_called_with(
         PipelineExecutionArn="my:arn"
     )
     assert len(steps) == 1
     list_parameters_response = execution.list_parameters()
-    assert (
-        sagemaker_session_mock.sagemaker_client.list_pipeline_parameters_for_execution.called_with(
-            PipelineExecutionArn="my:arn"
-        )
+    sagemaker_session_mock.sagemaker_client.list_pipeline_parameters_for_execution.assert_called_with(
+        PipelineExecutionArn="my:arn"
     )
+
     parameter_list = list_parameters_response["PipelineParameters"]
     assert len(parameter_list) == 1
     assert parameter_list[0]["Name"] == "TestParameterName"