def test_step_input_order_validation():
workflow_input = ExecutionInput()
test_step_01 = Pass(state_id='StateOne',
parameters={
'ParamA': workflow_input['Key02']['Key03'],
'ParamD': workflow_input['Key01']['Key03'],
})
test_step_02 = Pass(state_id='StateTwo',
parameters={
'ParamC': workflow_input["Key05"],
"ParamB": "SampleValueB",
"ParamE":
test_step_01.output()["Response"]["Key04"]
})
test_step_03 = Pass(state_id='StateThree',
parameters={
'ParamG': "SampleValueG",
"ParamF": workflow_input["Key06"],
"ParamH": "SampleValueH"
})
workflow_definition = Chain([test_step_01, test_step_03, test_step_02])
with pytest.raises(ValueError):
result = Graph(workflow_definition).to_dict()
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 __init__(self,
name,
definition,
role,
tags=[],
execution_input=None,
timeout_seconds=None,
comment=None,
version=None,
state_machine_arn=None,
format_json=True,
client=None):
"""
Args:
name (str): The name of the workflow. A name must not contain:
* whitespace
* brackets < > { } [ ]
* wildcard characters ? *
* special characters " # % \\ ^ | ~ ` $ & , ; : /
* control characters (U+0000-001F , U+007F-009F )
definition (str): The `Amazon States Language <https://states-language.net/spec.html>`_ definition of the workflow.
role (str): The Amazon Resource Name (ARN) of the IAM role to use for creating, managing, and running the workflow.
tags (list): Tags to be added when creating a workflow. Tags are key-value pairs that can be associated with Step Functions workflows and activities. (default: [])
execution_input (ExecutionInput, optional): Placeholder collection that defines the placeholder variables for the workflow execution. \
This is also used to validate inputs provided when executing the workflow. (default: None)
timeout_seconds (int, optional): The maximum number of seconds an execution of the workflow can run. If it runs longer than the specified time, the workflow run fails with a `States.Timeout` Error Name. (default: None)
comment (str, optional): A human-readable description of the workflow. (default: None)
version (str, optional): The version of the Amazon States Language used in the workflow. (default: None)
state_machine_arn (str, optional): The Amazon Resource Name (ARN) of the workflow. (default: None)
format_json (bool, optional): Boolean flag set to `True` if workflow definition and execution inputs should be prettified for this workflow. `False`, otherwise. (default: True)
client (SFN.Client, optional): boto3 client to use for creating, managing, and running the workflow on Step Functions. If not provided, a default boto3 client for Step Functions will be automatically created and used. (default: None)
"""
self.timeout_seconds = timeout_seconds
self.comment = comment
self.version = version
if isinstance(definition, Graph):
self.definition = definition
else:
self.definition = Graph(definition,
timeout_seconds=self.timeout_seconds,
comment=self.comment,
version=self.version)
self.name = name
self.role = role
self.tags = tags
self.workflow_input = execution_input
if client:
self.client = client
else:
self.client = boto3.client('stepfunctions')
append_user_agent_to_client(self.client)
self.format_json = format_json
self.state_machine_arn = state_machine_arn
def test_map_state_with_placeholders():
workflow_input = ExecutionInput()
step_result = StepResult()
map_state = Map(state_id="MapState01",
result_selector={
"foo": step_result["foo"],
"bar": step_result["bar1"]["bar2"]
})
iterator_state = Pass("TrainIterator",
parameters={
"ParamA": map_state.output()["X"]["Y"],
"ParamB":
workflow_input["Key01"]["Key02"]["Key03"]
})
map_state.attach_iterator(iterator_state)
workflow_definition = Chain([map_state])
expected_repr = {
"StartAt": "MapState01",
"States": {
"MapState01": {
"Type": "Map",
"ResultSelector": {
"foo.$": "$['foo']",
"bar.$": "$['bar1']['bar2']"
},
"End": True,
"Iterator": {
"StartAt": "TrainIterator",
"States": {
"TrainIterator": {
"Parameters": {
"ParamA.$":
"$['X']['Y']",
"ParamB.$":
"$$.Execution.Input['Key01']['Key02']['Key03']"
},
"Type": "Pass",
"End": True
}
}
}
}
}
}
result = Graph(workflow_definition).to_dict()
assert result == expected_repr
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_nested_parallel_example():
nested_level_1 = Parallel('NestedStateLevel1')
nested_level_1.add_branch(Succeed('NestedStateLevel2'))
first_state = Parallel('FirstState')
first_state.add_branch(nested_level_1)
result = Graph(first_state,
comment='This is a test.',
version='1.0',
timeout_seconds=3600).to_dict()
assert result == {
'StartAt': 'FirstState',
'Comment': 'This is a test.',
'Version': '1.0',
'TimeoutSeconds': 3600,
'States': {
'FirstState': {
'Type':
'Parallel',
'Branches': [{
'StartAt': 'NestedStateLevel1',
'States': {
'NestedStateLevel1': {
'Type':
'Parallel',
'Branches': [{
'StartAt': 'NestedStateLevel2',
'States': {
'NestedStateLevel2': {
'Type': 'Succeed'
}
}
}],
'End':
True
}
}
}],
'End':
True
}
}
}
def update(self, definition=None, role=None):
"""
Updates an existing state machine by modifying its definition and/or role. Executions started immediately after calling this method may use the previous definition and role.
Args:
definition (State or Chain, optional): The `Amazon States Language <https://states-language.net/spec.html>`_ definition to update the workflow with. (default: None)
role (str, optional): The Amazon Resource Name (ARN) of the IAM role to use for creating, managing, and running the workflow. (default: None)
Returns:
str: The state machine definition and/or role updated. If the update fails, None will be returned.
"""
if definition is None and role is None:
raise MissingRequiredParameter(
"A new definition and/or role must be provided to update an existing workflow."
)
if self.state_machine_arn is None:
raise WorkflowNotFound(
"Local workflow instance does not point to an existing workflow on AWS StepFunctions. Please consider using Workflow.create(...) to create a new workflow, or Workflow.attach(...) to attach the instance to an existing workflow on AWS Step Functions."
)
if definition:
if isinstance(definition, Graph):
self.definition = definition
else:
self.definition = Graph(definition,
timeout_seconds=self.timeout_seconds,
comment=self.comment,
version=self.version)
if role:
self.role = role
response = self.client.update_state_machine(
stateMachineArn=self.state_machine_arn,
definition=self.definition.to_json(pretty=self.format_json),
roleArn=self.role)
logger.info(
"Workflow updated successfully on AWS Step Functions. All execute() calls will use the updated definition and role within a few seconds. "
)
return self.state_machine_arn
def test_map_state_with_placeholders():
workflow_input = ExecutionInput()
map_state = Map('MapState01')
iterator_state = Pass('TrainIterator',
parameters={
'ParamA': map_state.output()['X']["Y"],
'ParamB':
workflow_input["Key01"]["Key02"]["Key03"]
})
map_state.attach_iterator(iterator_state)
workflow_definition = Chain([map_state])
expected_repr = {
"StartAt": "MapState01",
"States": {
"MapState01": {
"Type": "Map",
"End": True,
"Iterator": {
"StartAt": "TrainIterator",
"States": {
"TrainIterator": {
"Parameters": {
"ParamA.$":
"$['X']['Y']",
"ParamB.$":
"$$.Execution.Input['Key01']['Key02']['Key03']"
},
"Type": "Pass",
"End": True
}
}
}
}
}
}
result = Graph(workflow_definition).to_dict()
assert result == expected_repr
def test_workflow_with_placeholders():
workflow_input = ExecutionInput()
test_step_01 = Pass(state_id='StateOne',
parameters={
'ParamA': workflow_input['Key02']['Key03'],
'ParamD': workflow_input['Key01']['Key03'],
})
test_step_02 = Pass(state_id='StateTwo',
parameters={
'ParamC': workflow_input["Key05"],
"ParamB": "SampleValueB",
"ParamE":
test_step_01.output()["Response"]["Key04"]
})
test_step_03 = Pass(state_id='StateThree',
parameters={
'ParamG': "SampleValueG",
"ParamF": workflow_input["Key06"],
"ParamH": "SampleValueH"
})
workflow_definition = Chain([test_step_01, test_step_02, test_step_03])
result = Graph(workflow_definition).to_dict()
expected_workflow_repr = {
"StartAt": "StateOne",
"States": {
"StateOne": {
"Type": "Pass",
"Parameters": {
"ParamA.$": "$$.Execution.Input['Key02']['Key03']",
"ParamD.$": "$$.Execution.Input['Key01']['Key03']"
},
"Next": "StateTwo"
},
"StateTwo": {
"Type": "Pass",
"Parameters": {
"ParamC.$": "$$.Execution.Input['Key05']",
"ParamB": "SampleValueB",
"ParamE.$": "$['Response']['Key04']"
},
"Next": "StateThree"
},
"StateThree": {
"Type": "Pass",
"Parameters": {
"ParamG": "SampleValueG",
"ParamF.$": "$$.Execution.Input['Key06']",
"ParamH": "SampleValueH"
},
"End": True
}
}
}
assert result == expected_workflow_repr
def test_parallel_state_with_placeholders():
workflow_input = ExecutionInput()
parallel_state = Parallel('ParallelState01')
branch_A = Pass('Branch_A',
parameters={
'ParamA': parallel_state.output()['A']["B"],
'ParamB': workflow_input["Key01"]
})
branch_B = Pass('Branch_B',
parameters={
'ParamA':
"TestValue",
'ParamB':
parallel_state.output()["Response"]["Key"]["State"]
})
branch_C = Pass('Branch_C',
parameters={
'ParamA':
parallel_state.output()['A']["B"].get("C", float),
'ParamB': "HelloWorld"
})
parallel_state.add_branch(branch_A)
parallel_state.add_branch(branch_B)
parallel_state.add_branch(branch_C)
workflow_definition = Chain([parallel_state])
result = Graph(workflow_definition).to_dict()
expected_repr = {
"StartAt": "ParallelState01",
"States": {
"ParallelState01": {
"Type":
"Parallel",
"End":
True,
"Branches": [{
"StartAt": "Branch_A",
"States": {
"Branch_A": {
"Parameters": {
"ParamA.$": "$['A']['B']",
"ParamB.$": "$$.Execution.Input['Key01']"
},
"Type": "Pass",
"End": True
}
}
}, {
"StartAt": "Branch_B",
"States": {
"Branch_B": {
"Parameters": {
"ParamA": "TestValue",
"ParamB.$": "$['Response']['Key']['State']"
},
"Type": "Pass",
"End": True
}
}
}, {
"StartAt": "Branch_C",
"States": {
"Branch_C": {
"Parameters": {
"ParamA.$": "$['A']['B']['C']",
"ParamB": "HelloWorld"
},
"Type": "Pass",
"End": True
}
}
}]
}
}
}
assert result == expected_repr
def test_graph_from_string():
g = Graph(Chain([Pass('HelloWorld')]))
g1 = FrozenGraph.from_json(g.to_json())
assert isinstance(g1, Graph)
assert g.to_dict() == g1.to_dict()
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
}
}
}
def test_parallel_state_with_placeholders():
workflow_input = ExecutionInput()
step_result = StepResult()
parallel_state = Parallel(state_id="ParallelState01",
result_selector={
"foo": step_result["foo"],
"bar": step_result["bar1"]["bar2"]
})
branch_A = Pass("Branch_A",
parameters={
"ParamA": parallel_state.output()["A"]["B"],
"ParamB": workflow_input["Key01"]
})
branch_B = Pass("Branch_B",
parameters={
"ParamA":
"TestValue",
"ParamB":
parallel_state.output()["Response"]["Key"]["State"]
})
branch_C = Pass("Branch_C",
parameters={
"ParamA":
parallel_state.output()["A"]["B"].get("C", float),
"ParamB": "HelloWorld"
})
parallel_state.add_branch(branch_A)
parallel_state.add_branch(branch_B)
parallel_state.add_branch(branch_C)
workflow_definition = Chain([parallel_state])
result = Graph(workflow_definition).to_dict()
expected_repr = {
"StartAt": "ParallelState01",
"States": {
"ParallelState01": {
"Type":
"Parallel",
"ResultSelector": {
"foo.$": "$['foo']",
"bar.$": "$['bar1']['bar2']"
},
"End":
True,
"Branches": [{
"StartAt": "Branch_A",
"States": {
"Branch_A": {
"Parameters": {
"ParamA.$": "$['A']['B']",
"ParamB.$": "$$.Execution.Input['Key01']"
},
"Type": "Pass",
"End": True
}
}
}, {
"StartAt": "Branch_B",
"States": {
"Branch_B": {
"Parameters": {
"ParamA": "TestValue",
"ParamB.$": "$['Response']['Key']['State']"
},
"Type": "Pass",
"End": True
}
}
}, {
"StartAt": "Branch_C",
"States": {
"Branch_C": {
"Parameters": {
"ParamA.$": "$['A']['B']['C']",
"ParamB": "HelloWorld"
},
"Type": "Pass",
"End": True
}
}
}]
}
}
}
assert result == expected_repr