Skip to content

Commit 837b45e

Browse files
authored
Merge branch 'master' into fix-serverless-integ-test
2 parents 5f1131c + bcdfc6c commit 837b45e

File tree

2 files changed

+24
-1
lines changed

2 files changed

+24
-1
lines changed

src/sagemaker/workflow/callback_step.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ class CallbackOutput:
4949
"""
5050

5151
output_name: str = attr.ib(default=None)
52-
output_type: CallbackOutputTypeEnum = attr.ib(default=CallbackOutputTypeEnum.String.value)
52+
output_type: CallbackOutputTypeEnum = attr.ib(default=CallbackOutputTypeEnum.String)
5353

5454
def to_request(self) -> RequestType:
5555
"""Get the request structure for workflow service calls."""

tests/unit/sagemaker/workflow/test_callback_step.py

+23
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,29 @@ def test_callback_step():
5353
}
5454

5555

56+
def test_callback_step_default_values():
57+
param = ParameterInteger(name="MyInt")
58+
outputParam1 = CallbackOutput(output_name="output1")
59+
cb_step = CallbackStep(
60+
name="MyCallbackStep",
61+
depends_on=["TestStep"],
62+
sqs_queue_url="https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
63+
inputs={"arg1": "foo", "arg2": 5, "arg3": param},
64+
outputs=[outputParam1],
65+
)
66+
cb_step.add_depends_on(["SecondTestStep"])
67+
assert cb_step.to_request() == {
68+
"Name": "MyCallbackStep",
69+
"Type": "Callback",
70+
"DependsOn": ["TestStep", "SecondTestStep"],
71+
"SqsQueueUrl": "https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue",
72+
"OutputParameters": [
73+
{"OutputName": "output1", "OutputType": "String"},
74+
],
75+
"Arguments": {"arg1": "foo", "arg2": 5, "arg3": param},
76+
}
77+
78+
5679
def test_callback_step_output_expr():
5780
param = ParameterInteger(name="MyInt")
5881
outputParam1 = CallbackOutput(output_name="output1", output_type=CallbackOutputTypeEnum.String)

0 commit comments

Comments
 (0)