Testing Step Functions Locally

Have you built a Step Function with many steps, retries and end states - but you are left wondering, how do I test this masterpiece to ensure it's as wonderful as I think it is? Then you've come to the right place! Have a look at how we test Step Functions locally to give you more confidence in your work.

---

As you may have seen in our previous posts, we love Step Functions. It's great being able to build your Step Function in the console, see the payloads passing through your states and everything going green for you to say “Wooh! You’ve stepped through a Step Function successfully.”. But what if it didn’t, and it’s actually not doing what you expect, it’s going red and throwing useless errors or worse, it’s green but not giving you the response you want. What do you need? Tests!

What does AWS provide to help you test?

AWS itself provides some basic tools required for testing Step Functions - no they’re not the silver bullet in which you can just quickly write and run to test your Step Functions - but they certainly give you a jump start.

Step Functions Local documentation states:

AWS Step Functions Local is a downloadable version of Step Functions that lets you develop and test applications using a version of Step Functions running in your own development environment.

With Step Functions Local you can test locally or as part of a pipeline. You can test your flows, inputs, outputs, retries, back-offs and error states to ensure it performs as you expect.

Note:

Step Functions Local can sometimes be behind the Step Functions feature set. We have noticed when a new feature is implemented in Step Functions, the Step Functions Local container image may not be updated to include those features immediately. This is understandably not ideal - but you can keep an eye on the container here for new versions in which AWS are actively updating.

How to get it up and running

At Instil, we knew that we needed to run these tests as part of the pipeline but also run them locally when developing or investigating issues. AWS kindly provides some help with running the tests via the AWS CLI which is great, but we wanted to create these tests to last and have them run as part of our deployment pipeline. So we found this solution.

Here’s what you need:

1. AWS Step Functions Local (Docker Image)
2. Testcontainers package
3. AWS SDK package
4. Wait For Expect package

Step 1: Have a look at your Step Function

The Step Functions Workflow Studio is great for building out your Step Function in the console. It makes creating your Step Function user-friendly and makes visualising it super easy. Here we have an example Step Function.

It has a couple of lambdas, a choice state for checking the response of the first lambda and some success and failure paths. It has 4 flows which we would want to test if I can count correctly:

1. Get Item Price → “Item Price <= 100” → Success
2. Get Item Price → “Item Price > 100” → Ask for verification → Success
3. Get Item Price → Fail
4. Get Item Price → “Item Price > 100” → Ask for verification → Fail

Now we have an idea of what we want to test from our Step Function, we can get to work.

Step 2: Download your ASL file from the Step Function Workflow Studio

To use the Step Function Local container, we need our Step Function in ASL (Amazon States Language) which is AWS’ own language for defining Step Functions and their states. You can do this from the Step Function console by exporting the JSON definition.

Step 3: Get that Docker container spinning

You need the container up and running to be able to run the Step Function locally within it, we used testcontainers to spin up the short-lived container and have it ready for testing.

import {GenericContainer} from "testcontainers";

const awsStepFunctionsLocalContainer = await new GenericContainer("amazon/aws-stepfunctions-local")
    .withExposedPorts(8083)
    .withBindMount("your-path-to/MockConfigFile.json", "/home/MockConfigFile.json", "ro")
    .withEnv("SFN_MOCK_CONFIG", "/home/MockConfigFile.json")
    .start();

Note:

• Test Containers picks a random free port on the host machine and uses 8083 above to map it, so you don’t need to worry about clashes.
• MockConfigFile.json is the file we use for mocking how the AWS services respond in your Step Function test executions, we will come to how to create those in the next step!

Step 4: Create your MockConfigFile

The use of a mock config file is how we define the test cases, flows and responses of AWS service integrations within the Step Function. It makes up the meat of your Step Function testing journey and ultimately controls how detailed you want your tests to be.

The mock config is a JSON file which according to AWS’ own documentation includes:

• StateMachines - The fields of this object represent state machines configured to use mocked service integrations.
• MockedResponse - The fields of this object represent mocked responses for service integration calls.

Here’s what ours looks like as a finished product below. Make sure the names of the steps are identical to those named in the ASL file i.e “Get Item” in the test case is “Get Item” from the ASL file.

Note:

A great thing you can do in this file also detailed in the AWS documentation is to test the retry and backoff behaviour of some of your steps. For example, you could test that a lambda responds with an error on its first invocation, automatically retries and then returns successfully on its second invocation. Something like this is shown in the MockedGetItemAbove100 mocked response below.

{
  "StateMachines": {
    "ItemPriceChecker": {
      "TestCases": {
        "shouldSuccessfullyGetItemWithPriceBelow100": {
          "Get Item Price": "MockedGetItemBelow100"
        },
        "shouldSuccessfullyGetItemAndVerifyWithPriceEqualOrAbove100": {
          "Get Item Price": "MockedGetItemAbove100",
          "Ask for verification": "MockedAskForVerificationSuccess"
        },
        "shouldFailToGetItem": {
          "Get Item Price": "MockedGenericLambdaFailure"
        },
        "shouldFailToVerifyItemWithPriceEqualOrAbove100": {
          "Get Item Price": "MockedGetItemAbove100",
          "Ask for verification": "MockedGenericLambdaFailure"
        }
      }
    }
  },
  "MockedResponses": {
    "MockedGetItemBelow100": {
      "0": {
        "Return": {
          "StatusCode": 200,
          "Payload": {
            "StatusCode": 200,
            "itemPrice": 80
          }
        }
      }
    },
    "MockedGetItemAbove100": {
      "0": {
        "Throw": {
          "Error": "Lambda.TimeoutException",
          "Cause": "Lambda timed out."
        }
      },
      "1": {
        "Return": {
          "StatusCode": 200,
          "Payload": {
            "StatusCode": 200,
            "itemPrice": 100
          }
        }
      }
    },
    "MockedAskForVerificationSuccess": {
      "0": {
        "Return": {"StatusCode": 200}
      }
    },    
    "MockedGenericLambdaFailure": {
      "0": {
        "Throw": {
          "Error":"Lambda.GenericLambdaFailure",
          "Cause":"The lambda failed generically."
        }
      }
    }
  }
}

Step 5: Prepping the tests

So you have the Step Function and test cases ready, all you need now is to get them running. This first function will get the client for the Step Function Local container and allow you to run commands against it for testing the local version of your Step Function:

import {SFNClient} from "@aws-sdk/client-sfn";

const sfnLocalClient = new SFNClient({
  endpoint: `http://${awsStepFunctionsLocalContainer?.getHost()}:${awsStepFunctionsLocalContainer?.getMappedPort(8083)}`,
  region: "eu-west-2",
  credentials: {
    accessKeyId: "test",
    secretAccessKey: "test",
    sessionToken: "test"
  }
});

Important:

As you can see, we used “test” above for the credentials. This is to ensure the Step Function doesn’t interact with our actual deployed environment in AWS.

Step Functions Local allows you to run tests against actual deployed services (so feel free to do so for your case) but since we have mocked the services using MockConfigFile.json then we don’t want to do that. By using fake credentials then it just defaults to the mocked services from our file.

Next, create your local Step Function instance in the docker container using the client just created.

import {CreateStateMachineCommand} from "@aws-sdk/client-sfn";
import {readFileSync} from "fs";

const localStepFunction = await sfnLocalClient.send(
  new CreateStateMachineCommand({
    definition: readFileSync("your-path-to/ItemPriceCheckerAsl.json", "utf8"),
    name: "ItemPriceChecker",
    roleArn: undefined
  })
);

You can then start a Step Function execution for one of the test cases. This will run the Step Function in the container and use the mocked AWS service integrations defined in the MockConfigFile.json to determine the path it takes. Here is the function you can use, we have it wrapped here so it can be ran for each specific test case.

The stepFunctionInput is a JSON string of what you would be passing in to the Step Function. In our case for the ItemPriceChecker there is no input to the Step Function as the item price is retrieved in the first step - so the input can be anything e.g {}. Make sure for your own Step Function to pass in any input required or use {} similar to the example if no input is required.

import {StartExecutionCommand, StartExecutionCommandOutput} from "@aws-sdk/client-sfn";

async function startStepFunctionExecution(testName: string, stepFunctionInput: string): Promise<StartExecutionCommandOutput> {
  return await sfnLocalClient.send(
    new StartExecutionCommand({
      stateMachineArn: `${
        localStepFunction.stateMachineArn as string
      }#${testName}`,
      input: stepFunctionInput
    })
  );
}

Step 6: Finally some testing!

Now you have your running Step Function execution for a particular test case, we need to actually test it worked. This is where AWS isn’t super helpful, there is no provided API for interacting with the Step Function execution and determining how the Step Function handled your test data. So we had to make our own! Sort of.

Here’s an example using the Step Function execution from above:

import {GetExecutionHistoryCommand, GetExecutionHistoryCommandOutput, StartExecutionCommandOutput} from "@aws-sdk/client-sfn";
import waitFor from "wait-for-expect";

it("should successfully get item with price below 100", async () => {
  const stepFunctionInput = {};
  const expectedOutput = JSON.stringify({
    StatusCode: 200,    
    itemPrice: 80
  });

  // This runs the Step Function and returns the execution details using the function created earlier in the post
  const stepFunctionExecutionResult = await startStepFunctionExecution(
    "shouldSuccessfullyGetItemWithPriceBelow100",
    stepFunctionInput
  );

  // This checks the states to ensure the execution successfully completed with the correct output
  await thenTheItemPriceIsReturned(stepFunctionExecutionResult, expectedOutput);
});

async function thenTheItemPriceIsReturned(
  startLocalSFNExecutionResult: StartExecutionCommandOutput,
  expectedOutput: string
): Promise<void> {
  // Since the execution arn is provided, it could still be running so this waits for the execution to finish by checking for the result you need
  await waitFor(async () => {
    const getExecutionHistoryResult = await getExecutionHistory(startLocalSFNExecutionResult.executionArn);
    const successStateExitedEvent = getExecutionHistoryResult.events?.find(event => event.type === "SucceedStateExited");

    expect(successStateExitedEvent?.stateExitedEventDetails?.name).toEqual("Success");
    expect(successStateExitedEvent?.stateExitedEventDetails?.output).toEqual(expectedOutput);
  });
}

async function getExecutionHistory(executionArn: string | undefined): Promise<GetExecutionHistoryCommandOutput> {
  return await sfnLocalClient.send(
    new GetExecutionHistoryCommand({
      executionArn
    })
  );
}

There is a lot of information above but at its heart, it simply runs the Step Function in the container and returns the execution information to the test. It then grabs the execution history of the running local Step Function and checks for an event showing it succeeded; this allows the test to then also check the execution output and ensure it has succeeded correctly.

Step 7: Make sure to tear it all down

One thing that can be easily forgotten is your container running as a part of your test. A good thing to do is make sure it is torn down correctly at the end of your test run. This can be done very easily as part of an afterAll if running multiple tests and is simple done by stopping the test containers instance.

awsStepFunctionsLocalContainer.stop();

Step 8: Expand and add more tests

Now this is up to you! You can continue to test the rest of the flow cases for the Step Function, checking it has emitted “FailStateExited” in the execution history for the failed cases or expanding your testing flows.

The HistoryEventType from the aws-sdk gives you all the event types which can be logged in the Step Function Local execution history, this allows you to write tests however you like for checking the execution of the Step Function. Here are some examples of matcher functions we have written for different types of events:

import {HistoryEvent} from "@aws-sdk/client-sfn";

async findExecutionSucceededEventInHistory(executionArn: string | undefined): Promise<HistoryEvent | undefined> {
  return await findEventFromExecutionHistory(executionArn, "ExecutionSucceeded");
}

async findFailStateEnteredEventInHistory(executionArn: string | undefined): Promise<HistoryEvent | undefined> {
  return await findEventFromExecutionHistory(executionArn, "FailStateEntered");
}

async findSucceedStateExitedEventInHistory(executionArn: string | undefined): Promise<HistoryEvent | undefined> {
  return await findEventFromExecutionHistory(executionArn, "SucceedStateExited");
}

async findEventFromExecutionHistory(executionArn: string | undefined, eventKey: HistoryEventType): Promise<HistoryEvent | undefined> {
  const history = await getExecutionHistory(executionArn);

  return history.events?.find(
    event => event.type === eventKey
  );
}

You’re good to go!

What we have created above is hopefully something quite simple for testing Step Functions. We additionally improved this by creating a Step Function testing service class which holds all the re-usable functions and can be called easily within the required test file. With this we were able to run our Step Function tests as part of our deployment pipeline, providing greater confidence in our code and allowing us to integrate Step Functions more into our applications.

Important:

Now it's also good to note here that this is not everything we do at Instil to test our Step Functions, it is simply a companion that enables us to test the difficult edge cases including complicated flows, retries and back-offs etc. We are advocates for testing in the cloud - and this local testing mixed with integration testing in the cloud (focusing more on Step Functions interacting with other parts of the cloud rather than edge cases) is a good starting place for testing your Step Functions.

Additionally, we do hope to see some improvements to the Step Functions Local client in future from AWS, possibly providing their own matchers for checking that states have been entered and exited correctly within the tested Step Function, but if not we will just have to do it ourselves!