def test_find(make_stubber, found, error_code):
stepfunctions_client = boto3.client('stepfunctions')
stepfunctions_stubber = make_stubber(stepfunctions_client)
state_machine = StepFunctionsStateMachine(stepfunctions_client)
state_machine_name = 'test-state_machine_name'
state_machine_arn = 'test-arn'
machine_data = [('wrong-name', 'wrong-arn')]
if found:
machine_data.append((state_machine_name, state_machine_arn))
state_machines = [{
'name': name,
'stateMachineArn': arn
} for name, arn in machine_data]
stepfunctions_stubber.stub_list_state_machines(state_machines,
error_code=error_code)
if error_code is None:
got_state_machine_arn = state_machine.find(state_machine_name)
if found:
assert got_state_machine_arn == state_machine_arn
else:
assert got_state_machine_arn is None
else:
with pytest.raises(ClientError) as exc_info:
state_machine.find(state_machine_name)
assert exc_info.value.response['Error']['Code'] == error_code
def test_list_runs(make_stubber, run_status, error_code):
stepfunctions_client = boto3.client('stepfunctions')
stepfunctions_stubber = make_stubber(stepfunctions_client)
state_machine = StepFunctionsStateMachine(stepfunctions_client)
state_machine.state_machine_arn = 'test-arn'
runs = [{
'name': name,
'executionArn': arn
} for name, arn in [('run-name-1', 'run-arn-1'), ('run-name-2',
'run-arn-2')]]
stepfunctions_stubber.stub_list_executions(state_machine.state_machine_arn,
runs,
run_status,
error_code=error_code)
if error_code is None:
got_runs = state_machine.list_runs(run_status)
assert [{
'name': run['name'],
'executionArn': run['executionArn']
} for run in got_runs] == runs
else:
with pytest.raises(ClientError) as exc_info:
state_machine.list_runs(run_status)
assert exc_info.value.response['Error']['Code'] == error_code
def test_describe(make_stubber, error_code):
stepfunctions_client = boto3.client('stepfunctions')
stepfunctions_stubber = make_stubber(stepfunctions_client)
state_machine = StepFunctionsStateMachine(stepfunctions_client)
state_machine.state_machine_arn = 'test-state_machine_arn'
name = 'test-name'
definition = 'test-definition'
role_arn = 'test-role_arn'
stepfunctions_stubber.stub_describe_state_machine(
state_machine.state_machine_arn,
name,
definition,
role_arn,
error_code=error_code)
if error_code is None:
got_response = state_machine.describe()
assert got_response['name'] == name
assert got_response['definition'] == definition
assert got_response['roleArn'] == role_arn
assert got_response[
'stateMachineArn'] == state_machine.state_machine_arn
else:
with pytest.raises(ClientError) as exc_info:
state_machine.describe()
assert exc_info.value.response['Error']['Code'] == error_code
def test_create(make_stubber, error_code):
stepfunctions_client = boto3.client('stepfunctions')
stepfunctions_stubber = make_stubber(stepfunctions_client)
state_machine = StepFunctionsStateMachine(stepfunctions_client)
name = 'test-name'
definition = {'Comment': 'test-definition'}
role_arn = 'test-role-arn'
arn = 'test-arn'
stepfunctions_stubber.stub_create_state_machine(name,
definition,
role_arn,
arn,
error_code=error_code)
if error_code is None:
got_arn = state_machine.create(name, definition, role_arn)
assert got_arn == arn
assert state_machine.state_machine_name == name
assert state_machine.state_machine_arn == arn
else:
with pytest.raises(ClientError) as exc_info:
state_machine.create(name, definition, role_arn)
assert exc_info.value.response['Error']['Code'] == error_code
def test_delete(make_stubber, error_code):
stepfunctions_client = boto3.client('stepfunctions')
stepfunctions_stubber = make_stubber(stepfunctions_client)
state_machine = StepFunctionsStateMachine(stepfunctions_client)
state_machine.state_machine_arn = 'test-state_machine_arn'
stepfunctions_stubber.stub_delete_state_machine(
state_machine.state_machine_arn, error_code=error_code)
if error_code is None:
state_machine.delete()
else:
with pytest.raises(ClientError) as exc_info:
state_machine.delete()
assert exc_info.value.response['Error']['Code'] == error_code
def test_update(make_stubber, role_arn, error_code):
stepfunctions_client = boto3.client('stepfunctions')
stepfunctions_stubber = make_stubber(stepfunctions_client)
state_machine = StepFunctionsStateMachine(stepfunctions_client)
state_machine.state_machine_arn = 'test-arn'
definition = {'Comment': 'test-definition'}
stepfunctions_stubber.stub_update_state_machine(
state_machine.state_machine_arn,
definition,
role_arn,
error_code=error_code)
if error_code is None:
state_machine.update(definition, role_arn)
else:
with pytest.raises(ClientError) as exc_info:
state_machine.update(definition, role_arn)
assert exc_info.value.response['Error']['Code'] == error_code
def test_stop_run(make_stubber, error_code):
stepfunctions_client = boto3.client('stepfunctions')
stepfunctions_stubber = make_stubber(stepfunctions_client)
state_machine = StepFunctionsStateMachine(stepfunctions_client)
run_arn = 'test-run_arn'
cause = 'test cause'
stepfunctions_stubber.stub_stop_execution(run_arn,
cause,
error_code=error_code)
if error_code is None:
state_machine.stop_run(run_arn, cause)
else:
with pytest.raises(ClientError) as exc_info:
state_machine.stop_run(run_arn, cause)
assert exc_info.value.response['Error']['Code'] == error_code
def destroy(state_machine_name, stack, cf_resource):
"""
Destroys the state machine, the resources managed by the CloudFormation stack,
and the CloudFormation stack itself.
:param state_machine_name: The name of the state machine created for the demo.
:param stack: The CloudFormation stack that manages the demo resources.
:param cf_resource: A Boto3 CloudFormation resource.
"""
print("Removing the state machine.")
state_machine = StepFunctionsStateMachine(boto3.client('stepfunctions'))
state_machine.find(state_machine_name)
state_machine.delete()
print(f"Deleting {stack.name}.")
stack.delete()
print("Waiting for stack removal.")
waiter = cf_resource.meta.client.get_waiter('stack_delete_complete')
waiter.wait(StackName=stack.name)
print("Stack delete complete.")
def test_start_run(make_stubber, run_input, error_code):
stepfunctions_client = boto3.client('stepfunctions')
stepfunctions_stubber = make_stubber(stepfunctions_client)
state_machine = StepFunctionsStateMachine(stepfunctions_client)
state_machine.state_machine_arn = 'test-arn'
run_name = 'test-run_name'
run_arn = 'test-run_arn'
stepfunctions_stubber.stub_start_execution(state_machine.state_machine_arn,
run_name,
run_arn,
run_input,
error_code=error_code)
if error_code is None:
got_run_arn = state_machine.start_run(run_name, run_input)
assert got_run_arn == run_arn
else:
with pytest.raises(ClientError) as exc_info:
state_machine.start_run(run_name, run_input)
assert exc_info.value.response['Error']['Code'] == error_code
def usage_demo(state_machine_name, resources):
"""
Creates and runs a Step Functions state machine that calls a Lambda function to
retrieve message records from a DynamoDB table and record them as sent.
The state machine is then updated to also send messages to an Amazon SQS
queue and the state machine is run again.
"""
state_machine = StepFunctionsStateMachine(boto3.client('stepfunctions'))
table = boto3.resource('dynamodb').Table(resources['MessageTableName'])
queue = boto3.resource('sqs').Queue(resources['SendQueueUrl'])
state_machine_arn = state_machine.find(state_machine_name)
if state_machine_arn is None:
print("Create a message pump state machine.")
definition = make_definition(resources, False)
state_machine.create(state_machine_name, definition,
resources['StepRoleArn'])
print("Put three messages in the message table.")
for user_name, message in [
('wills', 'Brevity is the soul of wit.'),
('janea',
'Let us never underestimate the power of a well-written letter.'),
('lewisc',
"I have proved by actual trial that a letter, that takes an "
"hour to write, takes only about 3 minutes to read!")
]:
table.put_item(
Item={
'user_name': user_name,
'message': message,
'message_id': str(time.time_ns()),
'sent': False
})
print("Start the state machine.")
run_arn = state_machine.start_run(f"run-without-sqs-{time.time_ns()}")
print("Wait a few seconds for the state machine to run...")
time.sleep(10)
print("Verify that the messages in DynamoDB are marked as sent.")
messages = table.scan()['Items']
pprint(messages)
print("Stop the state machine.")
state_machine.stop_run(run_arn, "Stop to update for demo.")
runs = state_machine.list_runs('RUNNING')
while runs:
time.sleep(5)
runs = state_machine.list_runs('RUNNING')
print("Update the state machine so it sends messages to Amazon SQS.")
definition = make_definition(resources, True)
state_machine.update(definition)
time.sleep(5)
print("Reset the messages in the DynamoDB table to not sent.")
for msg in table.scan()['Items']:
table.update_item(Key={
'user_name': msg['user_name'],
'message_id': msg['message_id']
},
UpdateExpression='SET sent=:s',
ExpressionAttributeValues={':s': False})
print("Restart the state machine.")
run_arn = state_machine.start_run(f"run-with-sqs-{time.time_ns()}")
print("Wait for state machine to process messages...")
time.sleep(15)
print("Retrieve messages from Amazon SQS.")
poll_for_messages(queue)
print("Put another message in the table.")
table.put_item(
Item={
'user_name': 'wills',
'message': 'Action is eloquence.',
'message_id': str(time.time_ns()),
'sent': False
})
print("Give the state machine time to find and process the message.")
time.sleep(15)
print("Get messages from Amazon SQS.")
poll_for_messages(queue)
print("Stop the run.")
state_machine.stop_run(run_arn, "Done with demo.")