|
19 | 19 | import pytest
|
20 | 20 |
|
21 | 21 | from mock import Mock, call, patch
|
22 |
| -from mock.mock import MagicMock |
| 22 | +from mock.mock import MagicMock, ANY |
23 | 23 |
|
24 | 24 | from sagemaker import s3
|
25 | 25 | from sagemaker.remote_function.job import _JobSettings
|
@@ -345,6 +345,10 @@ def test_pipeline_update_with_parallelism_config(sagemaker_session_mock, role_ar
|
345 | 345 | role_arn=role_arn,
|
346 | 346 | parallelism_config=dict(MaxParallelExecutionSteps=10),
|
347 | 347 | )
|
| 348 | + pipeline.update( |
| 349 | + role_arn=role_arn, |
| 350 | + parallelism_config={"MaxParallelExecutionSteps": 10}, |
| 351 | + ) |
348 | 352 | sagemaker_session_mock.sagemaker_client.update_pipeline.assert_called_with(
|
349 | 353 | PipelineName="MyPipeline",
|
350 | 354 | PipelineDefinition=pipeline.definition(),
|
@@ -552,17 +556,17 @@ def test_pipeline_start(sagemaker_session_mock):
|
552 | 556 | sagemaker_session=sagemaker_session_mock,
|
553 | 557 | )
|
554 | 558 | pipeline.start()
|
555 |
| - sagemaker_session_mock.start_pipeline_execution.assert_called_with( |
| 559 | + sagemaker_session_mock.sagemaker_client.start_pipeline_execution.assert_called_with( |
556 | 560 | PipelineName="MyPipeline",
|
557 | 561 | )
|
558 | 562 |
|
559 | 563 | pipeline.start(execution_display_name="pipeline-execution")
|
560 |
| - sagemaker_session_mock.start_pipeline_execution.assert_called_with( |
| 564 | + sagemaker_session_mock.sagemaker_client.start_pipeline_execution.assert_called_with( |
561 | 565 | PipelineName="MyPipeline", PipelineExecutionDisplayName="pipeline-execution"
|
562 | 566 | )
|
563 | 567 |
|
564 | 568 | pipeline.start(parameters=dict(alpha="epsilon"))
|
565 |
| - sagemaker_session_mock.start_pipeline_execution.assert_called_with( |
| 569 | + sagemaker_session_mock.sagemaker_client.start_pipeline_execution.assert_called_with( |
566 | 570 | PipelineName="MyPipeline", PipelineParameters=[{"Name": "alpha", "Value": "epsilon"}]
|
567 | 571 | )
|
568 | 572 |
|
@@ -912,7 +916,7 @@ def test_pipeline_execution_basics(sagemaker_session_mock):
|
912 | 916 | PipelineExecutionArn="my:arn"
|
913 | 917 | )
|
914 | 918 | steps = execution.list_steps()
|
915 |
| - sagemaker_session_mock.sagemaker_client.describe_pipeline_execution_steps.assert_called_with( |
| 919 | + sagemaker_session_mock.sagemaker_client.list_pipeline_execution_steps.assert_called_with( |
916 | 920 | PipelineExecutionArn="my:arn"
|
917 | 921 | )
|
918 | 922 | assert len(steps) == 1
|
|
0 commit comments