Skip to content

Commit 08b10fb

Browse files
author
Rohan Gujarathi
committed
fix: flake8 error
1 parent 69dbd27 commit 08b10fb

File tree

2 files changed

+27
-6
lines changed

2 files changed

+27
-6
lines changed

src/sagemaker/workflow/lambda_step.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -82,8 +82,8 @@ def __init__(
8282
self,
8383
name: str,
8484
lambda_func: Lambda,
85-
inputs: dict = {},
86-
outputs: List[LambdaOutput] = [],
85+
inputs: dict = None,
86+
outputs: List[LambdaOutput] = None,
8787
cache_config: CacheConfig = None,
8888
depends_on: List[str] = None,
8989
):
@@ -103,15 +103,15 @@ def __init__(
103103
"""
104104
super(LambdaStep, self).__init__(name, StepTypeEnum.LAMBDA, depends_on)
105105
self.lambda_func = lambda_func
106-
self.outputs = outputs
106+
self.outputs = outputs if outputs is not None else []
107107
self.cache_config = cache_config
108-
self.inputs = inputs
108+
self.inputs = inputs if inputs is not None else {}
109109

110110
root_path = f"Steps.{name}"
111111
root_prop = Properties(path=root_path)
112112

113113
property_dict = {}
114-
for output in outputs:
114+
for output in self.outputs:
115115
property_dict[output.output_name] = Properties(
116116
f"{root_path}.OutputParameters['{output.output_name}']"
117117
)
@@ -150,7 +150,7 @@ def _get_function_arn(self):
150150
"""
151151
account_id = self.lambda_func.session.account_id()
152152
region = self.lambda_func.session.boto_region_name
153-
if "cn" in region.lower():
153+
if region.lower() == "cn-north-1" or region.lower() == "cn-northwest-1":
154154
partition = "aws-cn"
155155
else:
156156
partition = "aws"

tests/unit/sagemaker/workflow/test_lambda_step.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,3 +135,24 @@ def test_pipeline_interpolates_lambda_outputs():
135135
},
136136
],
137137
}
138+
139+
140+
def test_lambda_step_no_inputs_outputs():
141+
lambda_step = LambdaStep(
142+
name="MyLambdaStep",
143+
depends_on=["TestStep"],
144+
lambda_func=Lambda(
145+
function_arn="arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda"
146+
),
147+
inputs={},
148+
outputs=[],
149+
)
150+
lambda_step.add_depends_on(["SecondTestStep"])
151+
assert lambda_step.to_request() == {
152+
"Name": "MyLambdaStep",
153+
"Type": "Lambda",
154+
"DependsOn": ["TestStep", "SecondTestStep"],
155+
"FunctionArn": "arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda",
156+
"OutputParameters": [],
157+
"Arguments": {},
158+
}

0 commit comments

Comments
 (0)