Unit testing AWS: step function

Question

I am beginner in AWS and i have created my first AWS step function given below, Now next step is to unit test this step function. I independently unit tested my lambda function now i got stuck and have no idea about, how can i proceed for unit testing of step function.

I also get a question in my mind is it worth doing unit testing of step function,some time feel can it be done or not since it is just a json.

I tried to search but i didn't got any clue on internet or AWS documentation Any help will be appreciated any blog on this or any sample use case Thanks

{
"Comment": "An example of the Amazon States Language using a choice state.",
"StartAt": "LoginState",
States": {
"LoginState": {
  "Type": "Task",
  "Resource": "arn:aws:lambda:us-east-1:170344690019:function:myfirstLogin",
  "Next": "ChoiceState"
},
"ChoiceState": {
  "Type" : "Choice",
  "Choices": [
    {
      "Variable": "$.success",
      "BooleanEquals": true,
      "Next": "logoutState"
    },
    {
      "Variable": "$.success",
      "BooleanEquals": false,
      "Next": "DefaultState"
    }
  ],
  "Default": "DefaultState"
},

"logoutState": {
  "Type" : "Task",
  "Resource": "arn:aws:lambda:us-east-1:170344690019:function:myFirstLogout",
  "End": true
},


"DefaultState": {
  "Type": "Fail",
  "Error": "DefaultStateError",
  "Cause": "No Matches!"
}

}
}

Answer 1

This is a bit of a nitpick, but it will inform the following explanation. At the point of testing your state machine you are expanding outside the scope of unit testing into integration testing.

So why the nitpick? Since you are moving into integration testing you will need the ability to run the state machine so that you may feed it an input and validate the output. Here are 2 ways you can automate testing of your state machine...

1. Deploy your state machine into a test environment in your AWS account and invoke it directly using any of the tools provided by AWS (cli, boto3, etc.). This is closer to automation testing as it tests the state machine in a real environment. If you set this up as part of a CI pipeline it will require that you configure your build server with the access it needs to install and execute the state machine in your AWS account.

2. Try something like stepfunctions-local to emulate a running state machine on your local system or in your testing environment. This option could be useful if you have a CI pipeline setup that is already running your existing unit tests. This will take some effort to properly install the tooling into your CI environment but could be worth it.

3. My personal favorite...use localstack. These guys have done a great job of emulating several AWS services that be can brought up and run in a Docker container. This is particularly useful if your lambda uses other AWS services. I like to run it in my CI environment for integration testing.

4. Use AWS SAM CLI. I haven't used this much myself. It requires you to be using the serverless application model. Their documentation has really improved since it became more officially supported so it should be very easy to use by following their guides and numerous examples. Running this in a CI environment will require that you have the tool installed in your test environment.

I hope this helps. I don't think it would help to share any code in this answer because what you're trying to do isn't trivial and could be implemented multiple ways. For example, CI services like CircleCI leverage Docker containers giving you the option to generate your own Docker container for running stepfunctions-local or localstack.

EDIT

See the answer from @niqui below. I believe I would definitely favor this option for testing in a CI environment as an alternative to stepfunctions-local or localstack given that it is provided and maintained by AWS.

Answer 2

AWS announced recently a downloadable version of Step Functions

Answer 3

To mock the Lambda functions during interacting with the StepFunctions Local, a solution is creating a fake Lambda HTTP service in a Python thread initiated at the testing setup and making this service able to parse the HTTP request URL to determine which function to invoke.

I've implemented this concept as a pytest fixture: https://github.com/chehsunliu/pytest-stepfunctions.

Usage

Suppose there is a state machine which simply collects all the EMR cluster Ids and we want to test it locally.

State Machine Definition

{
  "StartAt": "ListIds",
  "States": {
    "ListIds": {
      "Type": "Task",
      "Resource": "${ListIdsLambdaArn}",
      "ResultPath": "$.cluster_ids",
      "End": true
    }
  }
}

Lambda Code

my/pkg/emr.py

import boto3


def list_ids(*args, **kwargs):
    emr_client = boto3.client("emr")
    response = emr_client.list_clusters()

    return [item["Id"] for item in response["Clusters"]]

Test Code

tests/test_foo.py

import json
import time
from string import Template

import boto3
from botocore.stub import Stubber


def test_bar(aws_stepfunctions_endpoint_url):
    # Create the definition string.

    definition_template = Template("""
    {
      "StartAt": "ListIds",
      "States": {
        "ListIds": {
          "Type": "Task",
          "Resource": "${ListIdsLambdaArn}",
          "ResultPath": "$.cluster_ids",
          "End": true
        }
      }
    }
    """)
    list_ids_lambda_arn = "arn:aws:lambda:us-east-1:123456789012:function:my.pkg.emr.list_ids"
    definition = definition_template.safe_substitute(ListIdsLambdaArn=list_ids_lambda_arn)

    # Create the state machine resource.

    sfn_client = boto3.client("stepfunctions", endpoint_url=aws_stepfunctions_endpoint_url)
    state_machine_arn = sfn_client.create_state_machine(
        name="list-ids", definition=definition, roleArn="arn:aws:iam::012345678901:role/DummyRole"
    )["stateMachineArn"]

    # Mock the Lambda code.

    emr_client = boto3.client("emr")
    mocker.patch("my.pkg.emr.boto3", autospec=True).client.return_value = emr_client

    stubber = Stubber(emr_client)
    stubber.add_response(
        "list_clusters", service_response={"Clusters": [{"Id": "j-00001"}, {"Id": "j-00002"}]}
    )

    # Start and wait until the execution finishes.

    execution_arn = sfn_client.start_execution(
        stateMachineArn=state_machine_arn, name="list-ids-exec", input="{}"
    )["executionArn"]

    with stubber:
        while True:
            response = sfn_client.describe_execution(executionArn=execution_arn)
            if response["status"] != "RUNNING":
                break
            time.sleep(0.5)

    # Validate the results.

    stubber.assert_no_pending_responses()
    assert "SUCCEEDED" == response["status"]
    assert ["j-00001", "j-00002"] == json.loads(response["output"])["cluster_ids"]

Running the test

Install the dependencies:

$ pip install boto3 pytest pytest-stepfunctions pytest-mock

Download the StepFunctions Local JAR here and execute it:

$ java -jar /path/to/StepFunctionsLocal.jar \
    --lambda-endpoint http://localhost:13000 \
    --step-functions-endpoint http://localhost:8083 \
    --wait-time-scale 0

Run the test:

$ python -m pytest -v \
    --pytest-stepfunctions-endpoint-url=http://0.0.0.0:8083 \
    --pytest-stepfunctions-lambda-address=0.0.0.0 \
    --pytest-stepfunctions-lambda-port=13000 \
    ./tests

The test can also be performed in Docker Compose, which is much easier to use and maintain. You can check the README in my repo. Hope this fixture could help people who found this article.

Answer 4

I had a similar problem, so I wrote an AWS unit tester for step functions. It works by using the official provided docker image.

Installation:

yarn add step-functions-tester
yarn add mocha chai

const TestRunner = require('step-functions-tester')
const { expect } = require('chai')
let testRunner
describe('Step function tester', function () {
  this.timeout('30s')

  before('Set up test runner', async function () {
    testRunner = new TestRunner()
    await testRunner.setUp()
  })
  afterEach('Clean up', async function () {
    await testRunner.cleanUp()
  })
  after('Tear down', async function () {
    await testRunner.tearDown()
  })

  it('Step function test', async function () {
    // AWS Step Function definition
    const stepFunctionDefinition = {StartAt: 'FirstStep', States: {FirstStep: { /* ... */}}}

    const stepFunctionInput = {}

    // Keys are function names in the step function definition, values are arrays of calls
    const callStubs = {'arn:eu-west:111:mockLambda': [{result: 'First call result'}, {result: 'Second call result'}], /*... */}
    
    const { executions } = await testRunner.run(callStubs, stepFunctionDefinition, stepFunctionInput)
    expect(executions).deep.equal(expectedExecutions)
  })
})