Skip to content

Commit 50ab9eb

Browse files
committed
fix test_pipeline
1 parent cbe7340 commit 50ab9eb

File tree

1 file changed

+9
-5
lines changed

1 file changed

+9
-5
lines changed

tests/unit/sagemaker/workflow/test_pipeline.py

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
import pytest
2020

2121
from mock import Mock, call, patch
22-
from mock.mock import MagicMock
22+
from mock.mock import MagicMock, ANY
2323

2424
from sagemaker import s3
2525
from sagemaker.remote_function.job import _JobSettings
@@ -345,6 +345,10 @@ def test_pipeline_update_with_parallelism_config(sagemaker_session_mock, role_ar
345345
role_arn=role_arn,
346346
parallelism_config=dict(MaxParallelExecutionSteps=10),
347347
)
348+
pipeline.update(
349+
role_arn=role_arn,
350+
parallelism_config={"MaxParallelExecutionSteps": 10},
351+
)
348352
sagemaker_session_mock.sagemaker_client.update_pipeline.assert_called_with(
349353
PipelineName="MyPipeline",
350354
PipelineDefinition=pipeline.definition(),
@@ -552,17 +556,17 @@ def test_pipeline_start(sagemaker_session_mock):
552556
sagemaker_session=sagemaker_session_mock,
553557
)
554558
pipeline.start()
555-
sagemaker_session_mock.start_pipeline_execution.assert_called_with(
559+
sagemaker_session_mock.sagemaker_client.start_pipeline_execution.assert_called_with(
556560
PipelineName="MyPipeline",
557561
)
558562

559563
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(
561565
PipelineName="MyPipeline", PipelineExecutionDisplayName="pipeline-execution"
562566
)
563567

564568
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(
566570
PipelineName="MyPipeline", PipelineParameters=[{"Name": "alpha", "Value": "epsilon"}]
567571
)
568572

@@ -912,7 +916,7 @@ def test_pipeline_execution_basics(sagemaker_session_mock):
912916
PipelineExecutionArn="my:arn"
913917
)
914918
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(
916920
PipelineExecutionArn="my:arn"
917921
)
918922
assert len(steps) == 1

0 commit comments

Comments
 (0)