def test_choice_state_creation():
choice_state = Choice('Choice', input_path='$.Input')
choice_state.add_choice(ChoiceRule.IsPresent("$.StringVariable1", True), Pass("End State 1"))
choice_state.add_choice(ChoiceRule.StringEquals("$.StringVariable1", "ABC"), Pass("End State 1"))
choice_state.add_choice(ChoiceRule.StringLessThanEqualsPath("$.StringVariable2", "$.value"), Pass("End State 2"))
choice_state.default_choice(Pass('End State 3'))
assert choice_state.state_id == 'Choice'
assert len(choice_state.choices) == 3
assert choice_state.default.state_id == 'End State 3'
assert choice_state.to_dict() == {
'Type': 'Choice',
'InputPath': '$.Input',
'Choices': [
{
'Variable': '$.StringVariable1',
'IsPresent': True,
'Next': 'End State 1'
},
{
'Variable': '$.StringVariable1',
'StringEquals': 'ABC',
'Next': 'End State 1'
},
{
'Variable': '$.StringVariable2',
'StringLessThanEqualsPath': '$.value',
'Next': 'End State 2'
}
],
'Default': 'End State 3'
}
with pytest.raises(TypeError):
Choice('Choice', unknown_field='Unknown Field')
def test_choice_state_with_placeholders():
first_state = Task(
'FirstState',
resource='arn:aws:lambda:us-east-1:1234567890:function:FirstState')
retry = Chain([Pass('Retry'), Pass('Cleanup'), first_state])
choice_state = Choice('Is Completed?')
choice_state.add_choice(
ChoiceRule.BooleanEquals(choice_state.output()["Completed"], True),
Succeed('Complete'))
choice_state.add_choice(
ChoiceRule.BooleanEquals(choice_state.output()["Completed"], False),
retry)
first_state.next(choice_state)
result = Graph(first_state).to_dict()
expected_repr = {
"StartAt": "FirstState",
"States": {
"FirstState": {
"Resource":
"arn:aws:lambda:us-east-1:1234567890:function:FirstState",
"Type": "Task",
"Next": "Is Completed?"
},
"Is Completed?": {
"Type":
"Choice",
"Choices": [{
"Variable": "$['Completed']",
"BooleanEquals": True,
"Next": "Complete"
}, {
"Variable": "$['Completed']",
"BooleanEquals": False,
"Next": "Retry"
}]
},
"Complete": {
"Type": "Succeed"
},
"Retry": {
"Type": "Pass",
"Next": "Cleanup"
},
"Cleanup": {
"Type": "Pass",
"Next": "FirstState"
}
}
}
assert result == expected_repr
def test_rule_serialization():
bool_rule = ChoiceRule.BooleanEquals('$.BooleanVariable', True)
assert bool_rule.to_dict() == {
'Variable': '$.BooleanVariable',
'BooleanEquals': True
}
string_rule = ChoiceRule.StringEquals('$.StringVariable', 'ABC')
assert string_rule.to_dict() == {
'Variable': '$.StringVariable',
'StringEquals': 'ABC'
}
and_rule = ChoiceRule.And([bool_rule, string_rule])
assert and_rule.to_dict() == {
'And': [{
'Variable': '$.BooleanVariable',
'BooleanEquals': True
}, {
'Variable': '$.StringVariable',
'StringEquals': 'ABC'
}]
}
not_rule = ChoiceRule.Not(string_rule)
assert not_rule.to_dict() == {
'Not': {
'Variable': '$.StringVariable',
'StringEquals': 'ABC'
}
}
compound_rule = ChoiceRule.Or([and_rule, not_rule])
assert compound_rule.to_dict() == {
'Or': [{
'And': [{
'Variable': '$.BooleanVariable',
'BooleanEquals': True
}, {
'Variable': '$.StringVariable',
'StringEquals': 'ABC'
}],
}, {
'Not': {
'Variable': '$.StringVariable',
'StringEquals': 'ABC'
}
}]
}
def test_wait_loop():
first_state = Task(
'FirstState',
resource='arn:aws:lambda:us-east-1:1234567890:function:FirstState')
retry = Chain([Pass('Retry'), Pass('Cleanup'), first_state])
choice_state = Choice('Is Completed?')
choice_state.add_choice(ChoiceRule.BooleanEquals('$.Completed', True),
Succeed('Complete'))
choice_state.add_choice(ChoiceRule.BooleanEquals('$.Completed', False),
retry)
first_state.next(choice_state)
result = Graph(first_state).to_dict()
assert result == {
'StartAt': 'FirstState',
'States': {
'FirstState': {
'Type': 'Task',
'Resource':
'arn:aws:lambda:us-east-1:1234567890:function:FirstState',
'Next': 'Is Completed?'
},
'Is Completed?': {
'Type':
'Choice',
'Choices': [{
'Variable': '$.Completed',
'BooleanEquals': True,
'Next': 'Complete'
}, {
'Variable': '$.Completed',
'BooleanEquals': False,
'Next': 'Retry'
}]
},
'Complete': {
'Type': 'Succeed'
},
'Retry': {
'Type': 'Pass',
'Next': 'Cleanup',
},
'Cleanup': {
'Type': 'Pass',
'Next': 'FirstState'
}
}
}
def test_variable_value_must_be_consistent():
string_functions = (
'StringEquals',
'StringLessThan',
'StringGreaterThan',
'StringLessThanEquals',
'StringGreaterThanEquals',
)
for string_function in string_functions:
func = getattr(ChoiceRule, string_function)
with pytest.raises(ValueError):
func('$.Variable', 42)
numeric_functions = (
'NumericEquals',
'NumericLessThan',
'NumericGreaterThan',
'NumericLessThanEquals',
'NumericGreaterThanEquals',
)
for numeric_function in numeric_functions:
func = getattr(ChoiceRule, numeric_function)
with pytest.raises(ValueError):
func('$.Variable', 'ABC')
with pytest.raises(ValueError):
ChoiceRule.BooleanEquals('$.Variable', 42)
timestamp_functions = (
'TimestampEquals',
'TimestampLessThan',
'TimestampGreaterThan',
'TimestampLessThanEquals',
'TimestampGreaterThanEquals',
)
for timestamp_function in timestamp_functions:
func = getattr(ChoiceRule, timestamp_function)
with pytest.raises(ValueError):
func('$.Variable', True)
crawl_database_check_status.add_retry(
sf.steps.Retry(error_equals=[
"Lambda.ServiceException", "Lambda.AWSLambdaException", "Lambda.SdkClientException"
])
)
crawl_database_check_status.add_catch(Catch(
error_equals=["States.TaskFailed"],
next_step=fail_state
))
crawl_database_check_choice = sf.steps.Choice(
state_id='Check finished?',
)
crawl_database_check_choice.add_choice(
ChoiceRule.BooleanEquals('$.crawlerStatus', False),
crawl_database_wait
)
crawl_database_check_choice.add_choice(
ChoiceRule.BooleanEquals('$.crawlerStatus', True),
finish
)
workflow_definition = sf.steps.Chain([
glue_import_to_s3_job,
crawl_params,
crawl_database_start,
crawl_database_wait,
crawl_database_check_status,
crawl_database_check_choice
])
def test_variable_must_start_with_prefix():
with pytest.raises(ValueError):
ChoiceRule.StringEquals('Variable', '42')
def test_string_matches_serialization():
string_matches_rule = ChoiceRule.StringMatches('$.input', 'hello*world\\*')
assert string_matches_rule.to_dict() == {
'Variable': '$.input',
'StringMatches': 'hello*world\\*'
}
'SAGEMAKER_AUTOPILOT_TARGET_MODEL',
'Type':
'PLAIN_TEXT',
'Value':
'{}.tar.gz'.format(execution_input['ModelName'])
}]
})
# happy path
model_and_endpoint_step = Chain(
[model_step, endpoint_config_step, endpoint_step, deploy_rest_api_task])
# define choice
check_job_choice.add_choice(
ChoiceRule.StringEquals(variable=check_autopilot_job_status.output()
['Payload']['AutopilotJobStatus'],
value='InProgress'),
next_step=check_autopilot_job_status)
check_job_choice.add_choice(
ChoiceRule.StringEquals(variable=check_autopilot_job_status.output()
['Payload']['AutopilotJobStatus'],
value='Stopping'),
next_step=check_autopilot_job_status)
check_job_choice.add_choice(
ChoiceRule.StringEquals(variable=check_autopilot_job_status.output()
['Payload']['AutopilotJobStatus'],
value='Failed'),
next_step=workflow_failure)
def test_choice_example():
next_state = Task(
'NextState',
resource='arn:aws:lambda:us-east-1:1234567890:function:NextState')
choice_state = Choice('ChoiceState')
choice_state.default_choice(
Fail('DefaultState', error='DefaultStateError', cause='No Matches!'))
choice_state.add_choice(
ChoiceRule.NumericEquals(variable='$.foo', value=1),
Chain([
Task('FirstMatchState',
resource=
'arn:aws:lambda:us-east-1:1234567890:function:FirstMatchState'
), next_state
]))
choice_state.add_choice(
ChoiceRule.NumericEquals(variable='$.foo', value=2),
Chain([
Task(
'SecondMatchState',
resource=
'arn:aws:lambda:us-east-1:1234567890:function:SecondMatchState'
), next_state
]))
chain = Chain()
chain.append(
Task(
'FirstState',
resource='arn:aws:lambda:us-east-1:1234567890:function:StartLambda'
))
chain.append(choice_state)
result = Graph(chain).to_dict()
assert result == {
'StartAt': 'FirstState',
'States': {
'FirstState': {
'Type': 'Task',
'Resource':
'arn:aws:lambda:us-east-1:1234567890:function:StartLambda',
'Next': 'ChoiceState'
},
'ChoiceState': {
'Type':
'Choice',
'Choices': [{
'Variable': '$.foo',
'NumericEquals': 1,
'Next': 'FirstMatchState'
}, {
'Variable': '$.foo',
'NumericEquals': 2,
'Next': 'SecondMatchState'
}],
'Default':
'DefaultState'
},
'FirstMatchState': {
'Type': 'Task',
'Resource':
'arn:aws:lambda:us-east-1:1234567890:function:FirstMatchState',
'Next': 'NextState'
},
'SecondMatchState': {
'Type': 'Task',
'Resource':
'arn:aws:lambda:us-east-1:1234567890:function:SecondMatchState',
'Next': 'NextState'
},
'DefaultState': {
'Type': 'Fail',
'Error': 'DefaultStateError',
'Cause': 'No Matches!'
},
'NextState': {
'Type': 'Task',
'Resource':
'arn:aws:lambda:us-east-1:1234567890:function:NextState',
'End': True
}
}
}
