def test_default_paths_not_included():
task_state = Task(
'Task',
resource='arn:aws:lambda:us-east-1:1234567890:function:StartLambda')
assert 'ResultPath' not in task_state.to_dict()
assert 'InputPath' not in task_state.to_dict()
assert 'OutputPath' not in task_state.to_dict()
def test_default_paths_not_converted_to_null():
task_state = Task(
'Task',
resource='arn:aws:lambda:us-east-1:1234567890:function:StartLambda')
assert '"ResultPath": null' not in task_state.to_json()
assert '"InputPath": null' not in task_state.to_json()
assert '"OutputPath": null' not in task_state.to_json()
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_wait_example():
chain = Chain()
chain.append(
Task('FirstState',
resource='arn:aws:lambda:us-east-1:1234567890:function:StartState'
))
chain.append(Wait('wait_using_seconds', seconds=10))
chain.append(Wait('wait_using_timestamp',
timestamp='2015-09-04T01:59:00Z'))
chain.append(
Wait('wait_using_timestamp_path', timestamp_path='$.expirydate'))
chain.append(
Wait('wait_using_seconds_path', seconds_path='$.expiryseconds'))
chain.append(
Task(
'FinalState',
resource='arn:aws:lambda:us-east-1:1234567890:function:EndLambda'))
result = Graph(chain).to_dict()
assert result == {
'StartAt': 'FirstState',
'States': {
'FirstState': {
'Type': 'Task',
'Resource':
'arn:aws:lambda:us-east-1:1234567890:function:StartState',
'Next': 'wait_using_seconds'
},
'wait_using_seconds': {
'Type': 'Wait',
'Seconds': 10,
'Next': 'wait_using_timestamp'
},
'wait_using_timestamp': {
'Type': 'Wait',
'Timestamp': '2015-09-04T01:59:00Z',
'Next': 'wait_using_timestamp_path'
},
'wait_using_timestamp_path': {
'Type': 'Wait',
'TimestampPath': '$.expirydate',
'Next': 'wait_using_seconds_path'
},
'wait_using_seconds_path': {
'Type': 'Wait',
'SecondsPath': '$.expiryseconds',
'Next': 'FinalState',
},
'FinalState': {
'Type': 'Task',
'Resource':
'arn:aws:lambda:us-east-1:1234567890:function:EndLambda',
'End': True
}
}
}
def test_task_state_creation():
task_state = Task('Task', resource='arn:aws:lambda:us-east-1:1234567890:function:StartLambda')
task_state.add_retry(Retry(error_equals=['ErrorA', 'ErrorB'], interval_seconds=1, max_attempts=2, backoff_rate=2))
task_state.add_retry(Retry(error_equals=['ErrorC'], interval_seconds=5))
task_state.add_catch(Catch(error_equals=['States.ALL'], next_step=Pass('End State')))
assert task_state.type == 'Task'
assert len(task_state.retries) == 2
assert len(task_state.catches) == 1
assert task_state.to_dict() == {
'Type': 'Task',
'Resource': 'arn:aws:lambda:us-east-1:1234567890:function:StartLambda',
'Retry': [
{
'ErrorEquals': ['ErrorA', 'ErrorB'],
'IntervalSeconds': 1,
'BackoffRate': 2,
'MaxAttempts': 2
},
{
'ErrorEquals': ['ErrorC'],
'IntervalSeconds': 5
}
],
'Catch': [
{
'ErrorEquals': ['States.ALL'],
'Next': 'End State'
}
],
'End': True
}
def test_task_state_creation_with_dynamic_timeout():
task_state = Task(
'Task',
resource='arn:aws:lambda:us-east-1:1234567890:function:StartLambda',
timeout_seconds_path='$.timeout',
heartbeat_seconds_path='$.heartbeat',
)
assert task_state.to_dict() == {
'Type': 'Task',
'Resource': 'arn:aws:lambda:us-east-1:1234567890:function:StartLambda',
'HeartbeatSecondsPath': '$.heartbeat',
'TimeoutSecondsPath': '$.timeout',
'End': True
}
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_task_state_create_fail_for_duplicated_dynamic_timeout_fields():
with pytest.raises(ValueError):
Task(
'Task',
resource='arn:aws:lambda:us-east-1:1234567890:function:StartLambda',
timeout_seconds=1,
timeout_seconds_path='$.timeout',
)
with pytest.raises(ValueError):
Task(
'Task',
resource='arn:aws:lambda:us-east-1:1234567890:function:StartLambda',
heartbeat_seconds=1,
heartbeat_seconds_path='$.heartbeat',
)
def test_step_result_placeholder():
step_result = StepResult()
test_step_01 = Task(state_id="StateOne",
result_selector={
"ParamA": step_result["foo"],
"ParamC": "SampleValueC"
})
expected_repr = {
"Type": "Task",
"ResultSelector": {
"ParamA.$": "$['foo']",
"ParamC": "SampleValueC"
},
"End": True
}
assert test_step_01.to_dict() == expected_repr
def test_paths_none():
task_state = Task('Task', resource='arn:aws:lambda:us-east-1:1234567890:function:StartLambda',
result_path=None,
input_path=None,
output_path=None)
assert 'ResultPath' in task_state.to_dict()
assert task_state.to_dict()['ResultPath'] is None
assert 'InputPath' in task_state.to_dict()
assert task_state.to_dict()['InputPath'] is None
assert 'OutputPath' in task_state.to_dict()
assert task_state.to_dict()['OutputPath'] is None
def test_schema_validation_for_step_result():
step_result = StepResult(schema={"Payload": {"Key01": str}})
with pytest.raises(ValueError):
test_step_01 = Task(state_id="StateOne",
result_selector={
"ParamA": step_result["Payload"]["Key02"],
"ParamB": "SampleValueB"
})
with pytest.raises(ValueError):
test_step_02 = Task(state_id="StateOne",
parameters={
"ParamA":
step_result["Payload"].get("Key01", float),
"ParamB":
"SampleValueB"
})
test_step_03 = Task(
state_id="StateOne",
result_selector={"ParamA": step_result["Payload"]["Key01"]})
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
}
}
}
'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',
endpoint_config_name=execution_input['EndpointConfigName'],
model_name=execution_input['ModelName'],
initial_instance_count=1,
instance_type='ml.m4.xlarge',
data_capture_config=DataCaptureConfig(
enable_capture=True,
sampling_percentage=100,
def __init__(self, datajob_stack, name):
super().__init__(datajob_stack, name)
self.sfn_task = Task(state_id="some-id")
def test_task_state_add_catch_adds_catcher_to_catchers(catch, expected_catch):
step = Task('Task')
step.add_catch(catch)
assert step.to_dict()['Catch'] == expected_catch
def test_task_state_constructor_with_catch_adds_catcher_to_catchers(
catch, expected_catch):
step = Task('Task', catch=catch)
assert step.to_dict()['Catch'] == expected_catch
def test_task_state_add_retry_adds_retrier_to_retriers(retry, expected_retry):
step = Task('Task')
step.add_retry(retry)
assert step.to_dict()['Retry'] == expected_retry
def test_task_state_constructor_with_retry_adds_retrier_to_retriers(
retry, expected_retry):
step = Task('Task', retry=retry)
assert step.to_dict()['Retry'] == expected_retry