def test_parallel_state_machine_creation(sfn_client, sfn_role_arn): parallel_state_name = "Parallel" left_pass_name = "Left Pass" right_pass_name = "Right Pass" final_state_name = "Final State" parallel_state_result = "Parallel Result" asl_state_machine_definition = { "StartAt": parallel_state_name, "States": { parallel_state_name: { "Type": "Parallel", "Next": final_state_name, "Branches": [ { "StartAt": left_pass_name, "States": { left_pass_name: { "Type": "Pass", "End": True } } }, { "StartAt": right_pass_name, "States": { right_pass_name: { "Type": "Pass", "End": True } } } ] }, final_state_name: { "Type": "Pass", "Result": parallel_state_result, "End": True } } } parallel_waits = steps.Parallel(parallel_state_name) parallel_waits.add_branch(steps.Pass(left_pass_name)) parallel_waits.add_branch(steps.Pass(right_pass_name)) definition = steps.Chain([ parallel_waits, steps.Pass(final_state_name, result=parallel_state_result) ]) workflow = Workflow( 'Test_Parallel_Workflow', definition=definition, role=sfn_role_arn ) workflow_test_suite(sfn_client, workflow, asl_state_machine_definition, parallel_state_result)
def test_map_state_machine_creation(sfn_client, sfn_role_arn): map_state_name = "Map State" iterated_state_name = "Pass State" final_state_name = "Final State" items_path = "$.array" max_concurrency = 0 map_state_result = "Map Result" state_machine_input = { "array": [1, 2, 3] } asl_state_machine_definition = { "StartAt": map_state_name, "States": { map_state_name: { "ItemsPath": items_path, "Iterator": { "StartAt": iterated_state_name, "States": { iterated_state_name: { "Type": "Pass", "End": True } } }, "MaxConcurrency": max_concurrency, "Type": "Map", "Next": final_state_name }, final_state_name: { "Type": "Pass", "Result": map_state_result, "End": True } } } map_state = steps.Map( map_state_name, items_path=items_path, iterator=steps.Pass(iterated_state_name), max_concurrency=max_concurrency) definition = steps.Chain([ map_state, steps.Pass(final_state_name, result=map_state_result) ]) workflow = Workflow( 'Test_Map_Workflow', definition=definition, role=sfn_role_arn ) workflow_test_suite(sfn_client, workflow, asl_state_machine_definition, map_state_result, state_machine_input)
def test_workflow_update(client, workflow): client.update_state_machine = MagicMock( return_value={'updateDate': datetime.now()}) new_definition = steps.Pass('HelloWorld') new_role = 'arn:aws:iam::1234567890:role/service-role/StepFunctionsRoleNew' assert workflow.update(definition=new_definition, role=new_role) == state_machine_arn
def test_workflow_update_when_statemachinearn_is_none(client): workflow = Workflow(name=state_machine_name, definition=definition, role=role_arn, client=client) new_definition = steps.Pass('HelloWorld') with pytest.raises(WorkflowNotFound): workflow.update(definition=new_definition)
def test_workflow_creation_failure_duplicate_state_ids(client): improper_definition = steps.Chain( [steps.Pass('HelloWorld'), steps.Succeed('HelloWorld')]) with pytest.raises(ValueError): workflow = Workflow(name=state_machine_name, definition=improper_definition, role=role_arn, client=client)
def test_catch_state_machine_creation(sfn_client, sfn_role_arn, training_job_parameters): catch_state_name = "TaskWithCatchState" custom_error = "CustomError" task_failed_error = "States.TaskFailed" all_fail_error = "States.ALL" custom_error_state_name = "Custom Error End" task_failed_state_name = "Task Failed End" all_error_state_name = "Catch All End" catch_state_result = "Catch Result" task_resource = "arn:aws:states:::sagemaker:createTrainingJob.sync" # change the parameters to cause task state to fail training_job_parameters["AlgorithmSpecification"]["TrainingImage"] = "not_an_image" asl_state_machine_definition = { "StartAt": catch_state_name, "States": { catch_state_name: { "Resource": task_resource, "Parameters": training_job_parameters, "Type": "Task", "End": True, "Catch": [ { "ErrorEquals": [ all_fail_error ], "Next": all_error_state_name } ] }, all_error_state_name: { "Type": "Pass", "Result": catch_state_result, "End": True } } } task = steps.Task( catch_state_name, parameters=training_job_parameters, resource=task_resource ) task.add_catch( steps.Catch( error_equals=[all_fail_error], next_step=steps.Pass(all_error_state_name, result=catch_state_result) ) ) workflow = Workflow( 'Test_Catch_Workflow', definition=task, role=sfn_role_arn ) workflow_test_suite(sfn_client, workflow, asl_state_machine_definition, catch_state_result)
def test_pass_state_machine_creation(sfn_client, sfn_role_arn): pass_state_name = "Pass" pass_state_result = "Pass Result" asl_state_machine_definition = { "StartAt": pass_state_name, "States": { pass_state_name: { "Result": pass_state_result, "Type": "Pass", "End": True } } } definition = steps.Pass(pass_state_name, result=pass_state_result) workflow = Workflow(unique_name_from_base('Test_Pass_Workflow'), definition=definition, role=sfn_role_arn) workflow_test_suite(sfn_client, workflow, asl_state_machine_definition, pass_state_result)
def test_task_state_machine_creation(sfn_client, sfn_role_arn, training_job_parameters): task_state_name = "TaskState" final_state_name = "FinalState" resource = "arn:aws:states:::sagemaker:createTrainingJob.sync" task_state_result = "Task State Result" asl_state_machine_definition = { "StartAt": task_state_name, "States": { task_state_name: { "Resource": resource, "Parameters": training_job_parameters, "Type": "Task", "Next": final_state_name }, final_state_name: { "Type": "Pass", "Result" : task_state_result, "End": True } } } definition = steps.Chain([ steps.Task( task_state_name, resource=resource, parameters=training_job_parameters ), steps.Pass(final_state_name, result=task_state_result) ]) workflow = Workflow( 'Test_Task_Workflow', definition=definition, role=sfn_role_arn ) workflow_test_suite(sfn_client, workflow, asl_state_machine_definition, task_state_result)
import uuid import boto3 import yaml import json from datetime import datetime from unittest.mock import MagicMock, Mock from stepfunctions import steps from stepfunctions.exceptions import WorkflowNotFound, MissingRequiredParameter from stepfunctions.workflow import Workflow, Execution, ExecutionStatus state_machine_name = 'HelloWorld' state_machine_arn = 'arn:aws:states:us-east-1:1234567890:stateMachine:HelloWorld' role_arn = 'arn:aws:iam::1234567890:role/service-role/StepFunctionsRole' execution_arn = 'arn:aws:states:us-east-1:1234567890:execution:HelloWorld:execution-1' definition = steps.Chain([steps.Pass('HelloWorld'), steps.Succeed('Complete')]) @pytest.fixture def client(): sfn = boto3.client('stepfunctions') sfn.describe_state_machine = MagicMock( return_value={ 'creationDate': datetime(2019, 9, 9, 9, 59, 59, 276000), 'definition': steps.Graph(definition).to_json(), 'name': state_machine_name, 'roleArn': role_arn, 'stateMachineArn': state_machine_arn, 'status': 'ACTIVE' }) sfn.create_state_machine = MagicMock(return_value={
def test_choice_state_machine_creation(sfn_client, sfn_role_arn): choice_state_name = "ChoiceState" first_match_name = "FirstMatchState" second_match_name = "SecondMatchState" default_state_name = "DefaultState" variable = "$.choice" first_choice_value = 1 second_choice_value = 2 default_error = "DefaultStateError" default_cause = "No Matches" first_choice_state_result = "First Choice State" second_choice_state_result = "Second Choice State" state_machine_input = {"choice": first_choice_value} asl_state_machine_definition = { "StartAt": choice_state_name, "States": { choice_state_name: { "Type": "Choice", "Choices": [{ "Variable": variable, "NumericEquals": first_choice_value, "Next": first_match_name }, { "Variable": variable, "NumericEquals": second_choice_value, "Next": second_match_name }], "Default": default_state_name }, default_state_name: { "Error": default_error, "Cause": default_cause, "Type": "Fail" }, first_match_name: { "Type": "Pass", "Result": first_choice_state_result, "End": True }, second_match_name: { "Type": "Pass", "Result": second_choice_state_result, "End": True } } } definition = steps.Choice(choice_state_name) definition.default_choice( steps.Fail(default_state_name, error=default_error, cause=default_cause)) definition.add_choice( steps.ChoiceRule.NumericEquals(variable=variable, value=first_choice_value), steps.Pass(first_match_name, result=first_choice_state_result)) definition.add_choice( steps.ChoiceRule.NumericEquals(variable=variable, value=second_choice_value), steps.Pass(second_match_name, result=second_choice_state_result)) workflow = Workflow(unique_name_from_base('Test_Choice_Workflow'), definition=definition, role=sfn_role_arn) workflow_test_suite(sfn_client, workflow, asl_state_machine_definition, first_choice_state_result, state_machine_input)
def test_wait_state_machine_creation(sfn_client, sfn_role_arn): first_state_name = "FirstState" first_wait_state_name = "WaitInSeconds" second_wait_state_name = "WaitTimestamp" third_wait_state_name = "WaitTimestampPath" fourth_wait_state_name = "WaitInSecondsPath" final_state_name = "FinalState" timestamp = "2019-09-04T01:59:00Z" timestamp_path = "$.expirydate" seconds = 2 seconds_path = "$.expiryseconds" wait_state_result = "Wait Result" parameters = {'expirydate': timestamp, 'expiryseconds': seconds} asl_state_machine_definition = { "StartAt": first_state_name, "States": { first_state_name: { "Type": "Pass", "Next": first_wait_state_name, "Parameters": parameters }, first_wait_state_name: { "Seconds": seconds, "Type": "Wait", "Next": second_wait_state_name }, second_wait_state_name: { "Timestamp": timestamp, "Type": "Wait", "Next": third_wait_state_name }, third_wait_state_name: { "TimestampPath": timestamp_path, "Type": "Wait", "Next": fourth_wait_state_name }, fourth_wait_state_name: { "SecondsPath": seconds_path, "Type": "Wait", "Next": final_state_name }, final_state_name: { "Type": "Pass", "Result": wait_state_result, "End": True } } } definition = steps.Chain([ steps.Pass(first_state_name, parameters=parameters), steps.Wait(first_wait_state_name, seconds=seconds), steps.Wait(second_wait_state_name, timestamp=timestamp), steps.Wait(third_wait_state_name, timestamp_path=timestamp_path), steps.Wait(fourth_wait_state_name, seconds_path=seconds_path), steps.Pass(final_state_name, result=wait_state_result) ]) workflow = Workflow(unique_name_from_base('Test_Wait_Workflow'), definition=definition, role=sfn_role_arn) workflow_test_suite(sfn_client, workflow, asl_state_machine_definition, wait_state_result)
def test_catch_state_machine_creation(sfn_client, sfn_role_arn, training_job_parameters): catch_state_name = "TaskWithCatchState" task_failed_error = "States.TaskFailed" timeout_error = "States.Timeout" task_failed_state_name = "Catch Task Failed End" timeout_state_name = "Catch Timeout End" catch_state_result = "Catch Result" task_resource = f"arn:{get_aws_partition()}:states:::sagemaker:createTrainingJob.sync" # Provide invalid TrainingImage to cause States.TaskFailed error training_job_parameters["AlgorithmSpecification"][ "TrainingImage"] = "not_an_image" task = steps.Task(catch_state_name, parameters=training_job_parameters, resource=task_resource, catch=steps.Catch(error_equals=[timeout_error], next_step=steps.Pass( timeout_state_name, result=catch_state_result))) task.add_catch( steps.Catch(error_equals=[task_failed_error], next_step=steps.Pass(task_failed_state_name, result=catch_state_result))) workflow = Workflow(unique_name_from_base('Test_Catch_Workflow'), definition=task, role=sfn_role_arn) asl_state_machine_definition = { "StartAt": catch_state_name, "States": { catch_state_name: { "Resource": task_resource, "Parameters": training_job_parameters, "Type": "Task", "End": True, "Catch": [{ "ErrorEquals": [timeout_error], "Next": timeout_state_name }, { "ErrorEquals": [task_failed_error], "Next": task_failed_state_name }] }, task_failed_state_name: { "Type": "Pass", "Result": catch_state_result, "End": True }, timeout_state_name: { "Type": "Pass", "Result": catch_state_result, "End": True }, } } workflow_test_suite(sfn_client, workflow, asl_state_machine_definition, catch_state_result)