def test_retry_fail_for_unsupported_state():
c1 = Choice('My Choice')
with pytest.raises(ValueError):
c1.add_catch(
Catch(error_equals=["States.NoChoiceMatched"],
next_step=Fail("ChoiceFailed")))
def test_choice_state_creation():
choice_state = Choice('Choice', input_path='$.Input')
choice_state.add_choice(ChoiceRule.StringEquals("$.StringVariable1", "ABC"), Pass("End State 1"))
choice_state.add_choice(ChoiceRule.StringLessThanEquals("$.StringVariable2", "ABC"), Pass("End State 2"))
choice_state.default_choice(Pass('End State 3'))
assert choice_state.state_id == 'Choice'
assert len(choice_state.choices) == 2
assert choice_state.default.state_id == 'End State 3'
assert choice_state.to_dict() == {
'Type': 'Choice',
'InputPath': '$.Input',
'Choices': [
{
'Variable': '$.StringVariable1',
'StringEquals': 'ABC',
'Next': 'End State 1'
},
{
'Variable': '$.StringVariable2',
'StringLessThanEquals': 'ABC',
'Next': 'End State 2'
}
],
'Default': 'End State 3'
}
with pytest.raises(TypeError):
Choice('Choice', unknown_field='Unknown Field')
def test_chaining_choice_with_existing_default_overrides_value(caplog):
s1_pass = Pass('Step - One')
s2_choice = Choice('Step - Two')
s3_pass = Pass('Step - Three')
s2_choice.default_choice(s3_pass)
# Chain s2_choice when default_choice is already set will trigger Warning message
with caplog.at_level(logging.WARNING):
Chain([s2_choice, s1_pass])
expected_warning = f'Chaining Choice state: Overwriting {s2_choice.state_id}\'s current default_choice ({s3_pass.state_id}) with {s1_pass.state_id}'
assert expected_warning in caplog.text
assert 'WARNING' in caplog.text
assert s2_choice.default == s1_pass
assert s2_choice.next_step is None # Choice steps do not have next_step
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_chaining_choice_sets_default_field():
s1_pass = Pass('Step - One')
s2_choice = Choice('Step - Two')
s3_pass = Pass('Step - Three')
chain1 = Chain([s1_pass, s2_choice, s3_pass])
assert chain1.steps == [s1_pass, s2_choice, s3_pass]
assert s1_pass.next_step == s2_choice
assert s2_choice.default == s3_pass
assert s2_choice.next_step is None # Choice steps do not have next_step
assert s3_pass.next_step is None
def test_append_states_after_terminal_state_will_fail():
with pytest.raises(ValueError):
chain = Chain()
chain.append(Pass('Pass'))
chain.append(Fail('Fail'))
chain.append(Pass('Pass2'))
with pytest.raises(ValueError):
chain = Chain()
chain.append(Pass('Pass'))
chain.append(Succeed('Succeed'))
chain.append(Pass('Pass2'))
with pytest.raises(ValueError):
chain = Chain()
chain.append(Pass('Pass'))
chain.append(Choice('Choice'))
chain.append(Pass('Pass2'))
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'
}
}
}
create_autopilot_job_step.add_catch(
Catch(error_equals=["States.TaskFailed"], next_step=workflow_failure))
check_autopilot_job_status = LambdaStep(
'CheckAutopilotJobStatus',
parameters={
'FunctionName': 'CheckAutopilotJobStatus',
'Payload': {
'AutopilotJobName':
create_autopilot_job_step.output()['Payload']['AutopilotJobName']
}
})
check_job_wait_state = Wait(state_id="Wait", seconds=360)
check_job_choice = Choice(state_id="IsAutopilotJobComplete")
model_step = Task(
'CreateAutopilotModel',
resource='arn:aws:states:::sagemaker:createModel',
parameters={
'Containers':
check_autopilot_job_status.output()['Payload']['InferenceContainers'],
'ModelName':
execution_input['ModelName'],
'ExecutionRoleArn':
sagemaker_exec_role
})
endpoint_config_step = EndpointConfigStep(
'CreateModelEndpointConfig',
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
}
}
}
