How to Test Step Functions State Machine Locally

Puneet Punj

27 March, 2023

Implement AWS Step Functions Local using Docker

AWS Step Functions is a serverless orchestration service that allows you to create and run state machines for coordinating distributed applications and microservices. It allows you to integrate with other AWS serverless tools such as AWS Lambda functions to build business-critical applications.

Fundamentally, AWS Step Functions operates as a simple state machine that defines the steps required for the execution of a specific workflow, and the methods that should be used to accomplish that. It allows you to build a multi-step workflow with parallel or serial execution of AWS services or simple Lambda functions, with an option to add choices and conditions.

There are a couple of options when it comes to building the state machine:

Use the AWS console and create a diagram of the entire process
Make use of Amazon States Language to create a json file, and deploy it using any Infrastructure as Code framework

All of this sounds great, but if you have ever worked with AWS Step Functions, you would know the most challenging part is testing the state machine locally. The good thing is that AWS has provided a great solution to test state machine logic locally using Step Functions Local. However, the documentation only offers basic HelloWorld implementation, and is not enough to implement in a real world state machine.

In this blog we will explore how to implement Step Functions Local using Docker to test all state machine paths in isolation, without any need to integrate with other AWS services.

Required Tools

Docker
Basic understanding of AWS Step Functions

Create Step Function

Let’s first create a basic state machine to stop applications running in an ECS cluster. You might need a similar state machine to stop a running application before starting the backup process. This step function will perform below tasks:

Invoke Lambda to stop all running ECS tasks
Invoke Lambda to get ECS status
The output of the above task will decide if state machine needs to wait for few seconds before checking status again
The state machine ends once the status check Lambda returns the ECS status as stopped

Create a file `statemachine/backup.json`

{
“Comment”: “State Machine to take application backup”,
“StartAt”: “ScaleInECSServices”,
“States”: {
“ScaleInECSServices”: {
“Type”: “Task”,
“Resource”: “arn:aws:states:::lambda:invoke”,
“OutputPath”: “$.Payload”,
“Parameters”: {
“Payload.$”: “$.ecs_cluster_name”,
“FunctionName”: “${scale_in_ecs_instance}”
},
“Next”: “CheckStatusECSServices”,
“Catch”: [
{
“ErrorEquals”: [“States.ALL”],
“Next”: “Final”
}
]
},
“CheckStatusECSServices”: {
“Type”: “Task”,
“Resource”: “arn:aws:states:::lambda:invoke”,
“OutputPath”: “$.Payload”,
“Parameters”: {
“Payload.$”: “$”,
“FunctionName”: “${check_status_ecs_services}”
},
“Next”: “AreAllServicesOffline?”,
“Catch”: [
{
“ErrorEquals”: [“States.ALL”],
“Next”: “Final”
}
]
},
“AreAllServicesOffline?”: {
“Type”: “Choice”,
“Choices”: [
{
“Variable”: “$.services_running”,
“BooleanEquals”: false,
“Next”: “Final”
},
{
“Variable”: “$.services_running”,
“BooleanEquals”: true,
“Next”: “WaitFor10Seconds”
}
]
},
“WaitFor10Seconds”: {
“Type”: “Wait”,
“Seconds”: 10,
“Next”: “CheckStatusECSServices”
},
“Final”: {
“Type”: “Pass”,
“End”: true
}
}
}

Set up Step Function Local

We are using two Lambda functions in this state machine. However, the focus of this article is on how to test different paths of the state machine. So, we will not consider the business logic of the individual Lambda function.

Step 1: Create `.env` file in the root of project with below details

AWS_ACCOUNT_ID=123456789013
AWS_DEFAULT_REGION=ap-southeast-2
SFN_MOCK_CONFIG=“/home/StepFunctionsLocal/MockConfigFile.json”

Step 2: Create `docker-compose.yml` file in the root of project with below details

services:
step_function_local:
container_name: “${CONTAINER_NAME-stepfunction_local}”
image: amazon/aws-stepfunctions-local
ports:
– “8083:8083”
env_file: .env
volumes:
– ./statemachine/test/MockConfigFile.json:/home/StepFunctionsLocal/MockConfigFile.json

Step 3: Create an event file `events/sfn_valid_input.json` to be used for testing

{
“ecs_cluster_name”: “test”
}

Create another event file `events/sfn_invalid_input.json` to be used for negative testing

{}

Step 4: Now, the main part of testing. Let’s create a file to write unit tests with mocked responses.

Create a new file `statemachine/test/MockConfigFile.json`

This is the file which Step Functions Local understands to run unit tests. It has below main keys:

“StateMachines” – Syntax to define which state machines are to be tested
“backup” – In this scenario, this is going to be the name of step function, but this could be any user defined name
“TestCases” – Syntax to define all the tests
“ECSServicesAreRunning”, “ECSServicesAreSopped” – These are our unit tests for the state machine created above
“MockedResponses” – Syntax key to start defining mocked responses for individual state machine task

{
“StateMachines”: {
“backup”: {
“TestCases”: {
“ECSServicesAreRunning”: {
“ScaleInECSServices”: “ScaleInECSLambdaMockedSuccess”,
“CheckStatusECSServices”: “CheckStatusECSServicesLambdaRunningTrue”
},
“ECSServicesAreStopped”: {
“ScaleInECSServices”: “ScaleInECSLambdaMockedSuccess”,
“CheckStatusECSServices”: “CheckStatusECSServicesLambdaRunningFalse”
},
“ScaleInServiceLambdaError”: {
“ScaleInECSServices”: “ScaleInECSServiceLambdaMockedThrowError”,
“CheckStatusECSServices”: “CheckStatusECSServicesLambdaRunningTrue”
},
“CheckStatusLambdaError”: {
“ScaleInECSServices”: “ScaleInECSLambdaMockedSuccess”,
“CheckStatusECSServices”: “CheckStatusLambdaMockedThrowError”
},
“InvalidInput”: {
“ScaleInECSServices”: “ScaleInECSLambdaMockedSuccess”,
“CheckStatusECSServices”: “CheckStatusECSServicesLambdaRunningTrue”
}
}
}
},
“MockedResponses”: {
“ScaleInECSLambdaMockedSuccess”: {
“0”: {
“Return”: {
“StatusCode”: 200,
“Payload”: {
“statusCode”: 200,
“body”: “{\”message\”: \”task count is set to 0\”}”
}
}
}
},
“CheckStatusECSServicesLambdaRunningTrue”: {
“0”: {
“Return”: {
“StatusCode”: 200,
“Payload”: {
“statusCode”: 200,
“services_running”: true
}
}
}
},
“CheckStatusECSServicesLambdaRunningFalse”: {
“0”: {
“Return”: {
“StatusCode”: 200,
“Payload”: {
“statusCode”: 200,
“services_running”: false
}
}
}
},
“ScaleInECSServiceLambdaMockedThrowError”: {
“0-3”: {
“Throw”: {
“Error”: “CustomValidationError”,
“Cause”: “ECS Cluster does not exist”
}
}
},
“CheckStatusLambdaMockedThrowError”: {
“0-3”: {
“Throw”: {
“Error”: “CustomValidationError”,
“Cause”: “Can not get ECS task status”
}
}
}
}
}

Step 5: Finally, let’s create a makefile in the root of the folder to have commands to build containers and run tests.
Create `makefile` and add the below targets to it.

ROOT_DIR:=$(shell dirname $(realpath $(firstword $(MAKEFILE_LIST))))
STATE_MACHINE_NAME=backup
STATE_MACHINE_DEFINITION_FILE=file://statemachine/${STATE_MACHINE_NAME}.json
STATE_MACHINE_ARN=arn:aws:states:ap-southeast-2:123456789013:stateMachine:${STATE_MACHINE_NAME}
STATE_MACHINE_EXECUTION_ARN=arn:aws:states:ap-southeast-2:123456789013:execution:${STATE_MACHINE_NAME}

run:
docker compose up -d –force-recreate

create:
aws stepfunctions create-state-machine \
–endpoint-url http://localhost:8083 \
–definition ${STATE_MACHINE_DEFINITION_FILE}\
–name ${STATE_MACHINE_NAME} \
–role-arn “arn:aws:iam::123456789012:role/DummyRole” \
–no-cli-pager

ecs_services_running:
aws stepfunctions start-execution \
–endpoint http://localhost:8083 \
–name ECSServicesAreRunningExecution \
–state-machine ${STATE_MACHINE_ARN}#ECSServicesAreRunning \
–input file://events/sfn_valid_input.json \
–no-cli-pager

ecs_services_stopped:
aws stepfunctions start-execution \
–endpoint http://localhost:8083 \
–name ECSServicesAreStoppedExecution \
–state-machine ${STATE_MACHINE_ARN}#ECSServicesAreStopped \
–input file://events/sfn_valid_input.json \
–no-cli-pager

scale_in_service_lambda_error:
aws stepfunctions start-execution \
–endpoint http://localhost:8083 \
–name ScaleInServiceErrorExecution \
–state-machine ${STATE_MACHINE_ARN}#ScaleInServiceLambdaError \
–input file://events/sfn_valid_input.json \
–no-cli-pager

check_status_lambda_error:
aws stepfunctions start-execution \
–endpoint http://localhost:8083 \
–name CheckStatusLambdaErrorExecution \
–state-machine ${STATE_MACHINE_ARN}#CheckStatusLambdaError \
–input file://events/sfn_valid_input.json \
–no-cli-pager

invalid_input:
aws stepfunctions start-execution \
–endpoint http://localhost:8083 \
–name InvalidInputExecution \
–state-machine ${STATE_MACHINE_ARN}#InvalidInput \
–input file://events/sfn_invalid_input.json \
–no-cli-pager

ecs_services_running_test:
aws stepfunctions get-execution-history \
–endpoint http://localhost:8083 \
–execution-arn ${STATE_MACHINE_EXECUTION_ARN}:ECSServicesAreRunningExecution \
–query ‘events[?type==`WaitStateEntered`]’ \
–no-cli-pager

ecs_services_stopped_test:
aws stepfunctions get-execution-history \
–endpoint http://localhost:8083 \
–execution-arn ${STATE_MACHINE_EXECUTION_ARN}:ECSServicesAreStoppedExecution \
–query ‘events[?type==`ExecutionSucceeded`]’ \
–no-cli-pager

scale_in_service_lambda_error_test:
aws stepfunctions get-execution-history \
–endpoint http://localhost:8083 \
–execution-arn ${STATE_MACHINE_EXECUTION_ARN}:ScaleInServiceErrorExecution \
–query ‘events[?type==`TaskFailed`]’ \
–no-cli-pager

check_status_lambda_error_test:
aws stepfunctions get-execution-history \
–endpoint http://localhost:8083 \
–execution-arn ${STATE_MACHINE_EXECUTION_ARN}:CheckStatusLambdaErrorExecution \
–query ‘events[?type==`TaskFailed`]’ \
–no-cli-pager

invalid_input_test:
aws stepfunctions get-execution-history \
–endpoint http://localhost:8083 \
–execution-arn ${STATE_MACHINE_EXECUTION_ARN}:InvalidInputExecution \
–query ‘events[?type==`ExecutionFailed`]’ \
–no-cli-pager

setup_all: ecs_services_running ecs_services_stopped scale_in_service_lambda_error check_status_lambda_error invalid_input

test_all: ecs_services_running_test ecs_services_stopped_test scale_in_service_lambda_error_test check_status_lambda_error_test invalid_input_test

Commands to execute tests locally

# set up docker container for aws step functions local
#this will allow us to use step functions locally on port 8083
make run

# create a new step function using endpoint http://localhost:8083
make create

#run command to start execution using event created above. here we specify which test to execute. this test was specified in the MockConfigFile.json (“ECSServicesAreRunning”, “ECSServicesAreSopped”, etc…)
make ecs_services_running

#finally run the step function get-exeuction-history command to test if we get a valid response based on mocks build in MockConfigFile.json
make ecs_services_running_test

#above command must respond with valid json instead of []

Framework Explanation

I know there is a lot going on above with different commands and files, so here is another attempt to visualise everything and explain how it all hangs together.

Mock configurations for your state machines
Name of the state machine under test
The number of test cases per the surrounding state machine that is under test
Name of the test case
Mapping of state (string match) with the supplied mock response
Mock responses used by all of the state machines under test
Mock response string matching the value from #5
Mock response for the first invocation of that state. Subsequent invocations can be referred to as “1”, “2”, and so on. In case of retries, it can be referred as “0-2” (for 3 failed retries) and another json object with key “3” mocking the successful response. This is assuming that the state machine state has Retry block with MaxAttempts set to 3
Return the mock response that matches the expected response from the task (response is not validated by Step Functions Local)

Makefile definitions:

M1 – Run command to start the step functions local Docker container. Re-run this command anytime you change MockConfigFIle.json
M2 – Command to create a new step function with a specified name in line #2
M2a – Use http://localhost:8083 endpoint to perform all step function operations
M2b – This is a step function definition file defined in line #3 and stored in the statemachine folder
M3 – This is a set up command to execute a step function. This will be successful unless there is an error in the actual step function definition. So this can help to prove the statemachine syntax is correct
- M3a – This is a name specified for step function execution. You could specify any user defined name here
- M3b – StateMachine ARN concatenated with #TestCaseName. This test case name is specified in the above file at 4th position. This has to match with the test case name mentioned in the MockConfigFile.json.
- M3c – This is an input event used to trigger step function
M4 – This is a command to get the step function execution history
- M4a – This must match with M3a
- M4b – This is an actual test case to check which event you are expecting after executing the step function. If the state machine does not work as per mocked responses or it has some other issue, this command will return an empty list []. Alternatively, a proper json will be returned and this is how you can prove this test is passed or not.

Conclusion

In this blog, we have stepped through how to develop and test Step Functions locally. However, there are couple of things you should keep in mind:

Step Functions Local can help you build step functions in isolation, however, will not help you to test and validate other integrated services. For that you would have to build additional solutions.
This approach can help you fail fast for the step functions logic

Github Repository – https://github.com/puneetpunj/step-function-local

Enjoyed this blog?

Share it with your network!

Migration

Data

Modernisation

AI

Launch

Co-Create

Modern Data Platform

Cevo AI Studio

Financial Services

Retail