def test_lambda_step(sagemaker_session):
param = ParameterInteger(name="MyInt")
outputParam1 = LambdaOutput(output_name="output1", output_type=LambdaOutputTypeEnum.String)
outputParam2 = LambdaOutput(output_name="output2", output_type=LambdaOutputTypeEnum.Boolean)
lambda_step = LambdaStep(
name="MyLambdaStep",
depends_on=["TestStep"],
lambda_func=Lambda(
function_arn="arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda",
session=sagemaker_session,
),
inputs={"arg1": "foo", "arg2": 5, "arg3": param},
outputs=[outputParam1, outputParam2],
)
lambda_step.add_depends_on(["SecondTestStep"])
assert lambda_step.to_request() == {
"Name": "MyLambdaStep",
"Type": "Lambda",
"DependsOn": ["TestStep", "SecondTestStep"],
"FunctionArn": "arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda",
"OutputParameters": [
{"OutputName": "output1", "OutputType": "String"},
{"OutputName": "output2", "OutputType": "Boolean"},
],
"Arguments": {"arg1": "foo", "arg2": 5, "arg3": param},
}
def test_lambda_step_no_inputs_outputs(sagemaker_session):
lambda_step = LambdaStep(
name="MyLambdaStep",
depends_on=["TestStep"],
lambda_func=Lambda(
function_arn="arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda",
session=sagemaker_session,
),
inputs={},
outputs=[],
)
lambda_step.add_depends_on(["SecondTestStep"])
assert lambda_step.to_request() == {
"Name": "MyLambdaStep",
"Type": "Lambda",
"DependsOn": ["TestStep", "SecondTestStep"],
"FunctionArn": "arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda",
"OutputParameters": [],
"Arguments": {},
}