Testing AWS Step Functions flows
Tags: aws, serverless, typescript, cdk
Published: 2022-02-14 17:00:47

---
title: Testing AWS Step Functions flows
published: true
description: Various techniques for testing AWS Step Functions flows.
tags: aws, serverless, typescript, cdk
cover_image: https://dev-to-uploads.s3.amazonaws.com/uploads/articles/hzn4dpynahqcnwwomn2u.jpeg
---
[AWS Step Functions](https://aws.amazon.com/step-functions/) is a serverless orchestrator service. Since its inception, it has become a go-to for all my low-code and orchestration needs on AWS.

Recently, AWS announced that the [Step Functions Local](https://docs.aws.amazon.com/step-functions/latest/dg/sfn-local.html) can now [mock service integrations](https://aws.amazon.com/blogs/compute/mocking-service-integrations-with-aws-step-functions-local/). I found the announcement a great opportunity and excuse to revisit the topic of testing in the context of AWS Step Functions, which historically was a bit hard.

This article will cover the techniques I use for testing flows that utilize the AWS Step Functions service as the logic orchestrator. Let us begin.

> The code in this blog post is written in [TypeScript](https://www.typescriptlang.org/) and uses [AWS CDK](https://aws.amazon.com/cdk/) for infrastructure piece. [This GitHub repository](https://github.com/WojciechMatuszewski/testing-step-functions) contains all the code used in this article.

{% github https://github.com/WojciechMatuszewski/testing-step-functions %}

## Layers of testing and confidence

Before we dive into specific techniques, we shall take a swift detour and talk about testing in general.

There are many heuristics when it comes to testing. There is the classical [testing pyramid](https://martinfowler.com/articles/practical-test-pyramid.html) or the [testing honeycomb](https://engineering.atspotify.com/2018/01/11/testing-of-microservices/) to name a few. In my personal opinion, **in the context of serverless applications**, tests written according to the testing honeycomb will give you much more confidence and ROI per test written than the "traditional" testing pyramid approach.

Before you base most of your tests on the local implementation of an AWS service, I urge you to think about the confidence the test gives you and less about how easy it is to write. Give the testing honeycomb a try. I'm positive you will not regret it!

## By writing integration / end-to-end tests

By utilizing e2e / integration tests, we gain the most confidence from our tests. That said, the tests are usually slow and can be hard to maintain to some degree.

Here, you will be interacting with AWS services directly, executing the AWS Step Function and asserting the side-effects of the flow. The side-effect can be, for example, a new item in [Amazon DynamoDB](https://aws.amazon.com/dynamodb/) or a message pushed to [Amazon SQS](https://aws.amazon.com/sqs/).

This style of testing is not free of challenging problems to tackle. Testing AWS Step Functions flow end-to-end becomes tricky when the Step Function definition is complex and contains multiple branches and error fallbacks.

Let us start with a basic example of saving a user into the Amazon DynamoDB table, then work our way up with Step Function definition complexity.

### Simplistic Step Functions workflows

The following image represents the Step Function definition we would like to test.

![Simplistic AWS Step Function definition](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/g27w90enjgffvll3lluv.png)

Lacking branching logic and error handling, all we have to do is test a single execution path.

> You should most likely always have error handling in place in your Step Function definition. This particular Step Function is only here for demonstration purposes. We will tackle testing error states later in the article.

Here is how I would write the test I'm referring to.

```ts
// simplistic-e2e.test.ts

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

const sfnClient = new SFNClient({});

test("Saves the user in the DynamoDB table", async () => {
  const startExecutionResult = await sfnClient.send(
    new StartExecutionCommand({
      stateMachineArn: process.env.SIMPLISTIC_E2E_STEP_FUNCTION_ARN,
      input: JSON.stringify({
        firstName: "John",
        lastName: "Doe"
      })
    })
  );

  await expect({
    region: process.env.AWS_REGION,
    table: process.env.SIMPLISTIC_E2E_DATA_TABLE_NAME
  }).toHaveItem(
    {
      PK: `USER#${startExecutionResult.executionArn}`
    },
    {
      PK: `USER#${startExecutionResult.executionArn}`,
      firstName: "John",
      lastName: "Doe"
    }
  );
}, 15_000);
```

First, I start the Step Function, and then I assert, using the [aws-testing-library](https://github.com/erezrokah/aws-testing-library), whether the item was correctly saved into the Amazon DynamoDB.

> Check out the [sls-test-tools](https://github.com/aleios-cloud/sls-test-tools) as well. It is a great library. I'm using _aws-testing-library_ because I'm used to it.

And that is it! Since this Step Function did not contain multiple branches and error handling, writing an end-to-end test takes little to no effort. With the most basic example behind us, let us tackle error states and multiple branches next.

### Step Functions with branching logic

Step Functions can be complex, especially when taking error handling and `Choice` steps into account. With the increased complexity comes increased difficulty in writing end-to-end tests.

Let us evolve our Step Function to include "background-check" [AWS Lambda](https://aws.amazon.com/lambda/) function. Depending on the result, the user in the DynamoDB table will have its `backgroundCheck` attribute populated (either "PASS", "FAIL" or "ERROR").

![Branching logic Step Function](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/dyhnbbvpzneruot4jnec.png)

> You can find the AWS Step Function AWS CDK definition [here](https://github.com/WojciechMatuszewski/testing-step-functions/blob/main/lib/branching-logic-e2e.ts).

To test all possible execution paths of the Step Function, **we have to have a way to force an error or particular response** for the `BackgroundCheckStep` AWS Lambda function – not an easy task!

So what can we do in this situation? Enter **aws-stepfunctions-local**.

The _aws-stepfunctions-local_ is a local implementation of the AWS Step Functions service provided to us by great folks at AWS. With this tool, we **will force the lambda to error without mocking other steps**.

> You could write a test that executes asserts on the status returned by the "background-check" AWS Lambda function. For the interest of time, I've chosen only to write a test for the case where the "background-check" fails.

Since I will be using the Docker version of the _aws-stepfunctions-local_, the **first step** is to integrate the act of spinning up and spinning down the container into the testing flow. My personal go-to in such situations is the [testcontainers](https://www.npmjs.com/package/testcontainers) package.

```ts
// branching-logic-e2e.test.ts

let container: StartedTestContainer | undefined;

beforeAll(async () => {
  const mockConfigPath = join(__dirname, "./branching-logic-e2e.mocks.json");
  container = await new GenericContainer("amazon/aws-stepfunctions-local")
    .withExposedPorts(8083)
    .withBindMount(mockConfigPath, "/home/branching-logic-e2e.mocks.json", "ro")
    .withEnv("SFN_MOCK_CONFIG", "/home/branching-logic-e2e.mocks.json")
    .withEnv("AWS_ACCESS_KEY_ID", process.env.AWS_ACCESS_KEY_ID as string)
    .withEnv(
      "AWS_SECRET_ACCESS_KEY",
      process.env.AWS_SECRET_ACCESS_KEY as string
    )
    // For federated credentials (for example, SSO), this environment variable is required.
    .withEnv("AWS_SESSION_TOKEN", process.env.AWS_SESSION_TOKEN as string)
    .withEnv("AWS_DEFAULT_REGION", process.env.AWS_REGION)
    .start();
}, 15_000);

afterAll(async () => {
  await container?.stop();
}, 15_000);
```

Let us unpack what is going on here.

First, the mysterious `mockConfigPath` and subsequent usages of this variable. The path points to a mock configuration file described on [_aws-stepfunctions-local_ documentation page](https://docs.aws.amazon.com/step-functions/latest/dg/sfn-local-test-sm-exec.html) and contains a mock definition for the `BackgroundCheckStep`.

```json
{
  "StateMachines": {
    "BranchingLogic": {
      "TestCases": {
        "ErrorPath": {
          "BackgroundCheckStep": "BackgroundCheckError"
        }
      }
    }
  },
  "MockedResponses": {
    "BackgroundCheckError": {
      "0": {
        "Throw": {
          "Error": "Lambda.TimeoutException",
          "Cause": "Lambda timed out."
        }
      }
    }
  }
}
```

Notice that I'm only concerned with the `BackgroundCheckStep` here. **By only mocking `BackgroundCheckStep`, I can force this particular step to fail**. Other steps are not mocked. Thus **the _aws-stepfunctions-local_ will reach out to native AWS services** – in our case Amazon DynamoDB.

The Docker image needs to have AWS-related environment variables populated to allow _aws-stepfunctions-local_ to talk to other AWS services. Refer to [this documentation page](https://docs.aws.amazon.com/step-functions/latest/dg/sfn-local-config-options.html#docker-credentials) for more information.

```ts
.withEnv("AWS_ACCESS_KEY_ID", process.env.AWS_ACCESS_KEY_ID as string)
.withEnv("AWS_SECRET_ACCESS_KEY",process.env.AWS_SECRET_ACCESS_KEY as string)
.withEnv("AWS_SESSION_TOKEN", process.env.AWS_SESSION_TOKEN as string)
.withEnv("AWS_DEFAULT_REGION", process.env.AWS_REGION)
```

As for the test itself, the first thing to do is gather necessary information about the Step Function under test.

```ts
// branching-logic-e2e.test.ts

test("Handles the failure of the BackgroundCheck step", async () => {
  const sfnClient = new SFNClient({});

  const describeStepFunctionResult = await sfnClient.send(
    new DescribeStateMachineCommand({
      stateMachineArn: process.env.BRANCHING_LOGIC_E2E_STEP_FUNCTION_ARN
    })
  );
  const stepFunctionDefinition =
    describeStepFunctionResult.definition as string;
  const stepFunctionRoleARN = describeStepFunctionResult.roleArn as string;

  // Rest of the test...
}, 50_000);
```

We will need all of this information since we will be re-creating this Step Function locally. Remember, the _aws-stepfunctions-local_ is the engine that runs our Step Function.

Next, let us re-create and run the Step Function that lives in the cloud using the _aws-stepfunctions-local_.

```ts
// branching-logic-e2e.test.ts

// Test setup ...

test("Handles the failure of the BackgroundCheck step", async () => {
  // Previous code snippet ...

  const sfnLocalClient = new SFNClient({
    endpoint: `http://localhost:${container?.getMappedPort(8083)}`
  });

  const createLocalSFNResult = await sfnLocalClient.send(
    new CreateStateMachineCommand({
      definition: stepFunctionDefinition,
      name: "BranchingLogic",
      roleArn: stepFunctionRoleARN
    })
  );

  const startLocalSFNExecutionResult = await sfnLocalClient.send(
    new StartExecutionCommand({
      stateMachineArn: `${
        createLocalSFNResult.stateMachineArn as string
      }#ErrorPath`,
      input: JSON.stringify({
        firstName: "John",
        lastName: "Doe"
      })
    })
  );

  // Rest of the test...
}, 50_000);
```

There are two **essential** pieces of detail I would like to bring your attention to.

1. The `name` of the Step Function I'm creating. The **`name` parameter must be the same as declared in the mock configuration file**.

2. The `stateMachineArn` format in the `StartExecutionCommand` call. The convention of `ARN#TEST_CASE` is required by _aws-stepfunctions-local_. Refer to [this documentation page](https://docs.aws.amazon.com/step-functions/latest/dg/sfn-local-test-sm-exec.html#run-mocked-serv-integ-tests) to learn more.

All that is left is to assert on the data of a given Amazon DynamoDB item. The assertion is almost identical as in the _Simplistic Step Functions_ section.

```ts
// branching-logic-e2e.test.ts

// Test setup ...

test("Handles the failure of the BackgroundCheck step", async () => {
  // Previous code snippet ...

  await expect({
    region: process.env.AWS_REGION,
    table: process.env.BRANCHING_LOGIC_E2E_DATA_TABLE_NAME
  }).toHaveItem(
    {
      PK: `USER#${startLocalSFNExecutionResult.executionArn}`
    },
    {
      PK: `USER#${startLocalSFNExecutionResult.executionArn}`,
      firstName: "John",
      lastName: "Doe",
      // The `BackgroundCheckStep` must have failed for this attribute to have the `ERROR` value.
      backgroundCheck: "ERROR"
    }
  );
}, 50_000);
```

We assert on an Amazon DynamoDB table that lives in the AWS. Since we **did not** provide any mocks for the step that saves the user data to Amazon DynamoDB, the _aws-stepfunctions-local_ made the request to the actual service, utilizing the credentials from Docker environment variables.

{% details Click to expand (the whole test definition) %}

```ts
import {
  CreateStateMachineCommand,
  DescribeStateMachineCommand,
  SFNClient,
  StartExecutionCommand
} from "@aws-sdk/client-sfn";
import { join } from "path";
import { GenericContainer, StartedTestContainer } from "testcontainers";

let container: StartedTestContainer | undefined;

beforeAll(async () => {
  const mockConfigPath = join(__dirname, "./branching-logic-e2e.mocks.json");
  container = await new GenericContainer("amazon/aws-stepfunctions-local")
    .withExposedPorts(8083)
    .withBindMount(mockConfigPath, "/home/branching-logic-e2e.mocks.json", "ro")
    .withEnv("SFN_MOCK_CONFIG", "/home/branching-logic-e2e.mocks.json")
    .withEnv("AWS_ACCESS_KEY_ID", process.env.AWS_ACCESS_KEY_ID as string)
    .withEnv(
      "AWS_SECRET_ACCESS_KEY",
      process.env.AWS_SECRET_ACCESS_KEY as string
    )
    /**
     * For federated credentials (for example, SSO), this environment variable is required.
     */
    .withEnv("AWS_SESSION_TOKEN", process.env.AWS_SESSION_TOKEN as string)
    .withEnv("AWS_DEFAULT_REGION", process.env.AWS_REGION)
    .start();
}, 15_000);

afterAll(async () => {
  await container?.stop();
}, 15_000);

test("Handles the failure of the BackgroundCheck step", async () => {
  const sfnClient = new SFNClient({});

  const describeStepFunctionResult = await sfnClient.send(
    new DescribeStateMachineCommand({
      stateMachineArn: process.env.BRANCHING_LOGIC_E2E_STEP_FUNCTION_ARN
    })
  );
  const stepFunctionDefinition =
    describeStepFunctionResult.definition as string;
  const stepFunctionRoleARN = describeStepFunctionResult.roleArn as string;

  const sfnLocalClient = new SFNClient({
    endpoint: `http://localhost:${container?.getMappedPort(8083)}`
  });

  const createLocalSFNResult = await sfnLocalClient.send(
    new CreateStateMachineCommand({
      definition: stepFunctionDefinition,
      name: "BranchingLogic",
      roleArn: stepFunctionRoleARN
    })
  );

  const startLocalSFNExecutionResult = await sfnLocalClient.send(
    new StartExecutionCommand({
      stateMachineArn: `${
        createLocalSFNResult.stateMachineArn as string
      }#ErrorPath`,
      input: JSON.stringify({
        firstName: "John",
        lastName: "Doe"
      })
    })
  );

  await expect({
    region: process.env.AWS_REGION,
    table: process.env.BRANCHING_LOGIC_E2E_DATA_TABLE_NAME
  }).toHaveItem(
    {
      PK: `USER#${startLocalSFNExecutionResult.executionArn}`
    },
    {
      PK: `USER#${startLocalSFNExecutionResult.executionArn}`,
      firstName: "John",
      lastName: "Doe",
      backgroundCheck: "ERROR"
    }
  );
}, 50_000);
```

{% enddetails %}

## By decomposing AWS Step Function Tasks

Switching gears from a somewhat complex end-to-end tests world, there exists another technique I would like to highlight.

I first saw this method of testing various AWS Step Function tasks while browsing code written by my colleagues at [Stedi](http://stedi.com/).

My friend [Graham Allan](https://twitter.com/Grundlefleck) wrote about it [on his blog here](https://grundlefleck.github.io/2022/01/12/how-using-the-same-programming-language-for-iac-made-a-step-function-testable.html). Check it out!

## Testing Pass state transformations

Have you ever spent way too much time trying to transform data via the `Pass` state to the shape you need it to be? I know I have.

As great as the AWS console and the [AWS Step Functions data flow simulator](https://aws.amazon.com/about-aws/whats-new/2021/04/aws-step-functions-adds-new-data-flow-simulator-for-modelling-input-and-output-processing/) is, I find the feedback loop of a test case unbeatable.

So, how can we reliably extract those `Pass` states from our AWS CDK code and test them? Here is my solution.

Here is our sample [_Construct_](https://docs.aws.amazon.com/cdk/v2/guide/constructs.html) definition. It contains the `transformDataStep` that we would like to test.

```ts
// pass-states-integration.ts

export class PassStatesIntegration extends Construct {
  constructor(scope: Construct, id: string) {
    super(scope, id);

    const transformDataStep = new aws_stepfunctions.Pass(
      this,
      "TransformDataStep",
      {
        parameters: {
          payload: aws_stepfunctions.JsonPath.stringAt(
            "States.Format('{} {}', $.firstName, $.lastName)"
          )
        }
      }
    );

    // You can imagine the definition being a bit more complex.
    const stepFunctionDefinition = transformDataStep;

    const stepFunction = new aws_stepfunctions.StateMachine(
      this,
      "StepFunction",
      {
        definition: stepFunctionDefinition
      }
    );
  }
}
```

We will be using the _aws-stepfunctions-local_, so the test setup looks very similar to the previous code snippets in this article.

```ts
// pass-states-integration.test.ts

let container: StartedTestContainer | undefined;

beforeAll(async () => {
  container = await new GenericContainer("amazon/aws-stepfunctions-local")
    .withExposedPorts(8083)
    .withEnv("AWS_DEFAULT_REGION", process.env.AWS_REGION)
    .start();
}, 15_000);

afterAll(async () => {
  await container?.stop();
}, 15_000);
```

To test the `transformDataStep`, we must retrieve the `transformDataStep` [ASL](https://docs.aws.amazon.com/step-functions/latest/dg/concepts-amazon-states-language.html) definition. We can do this in two ways.

1. Use the `DescribeStateMachine` API call, like we did in the _Step Functions with branching logic_ section.

2. Make the `transformDataStep` a publicly accessible property on the `PassStatesIntegration` construct.

I'm going to go with option number two since we have already seen option number one in action.

```diff
// pass-states-integration.ts

export class PassStatesIntegration extends Construct {
+ public transformDataStep: aws_stepfunctions.Pass;

  constructor(scope: Construct, id: string) {
    super(scope, id);

-   const transformDataStep = new aws_stepfunctions.Pass(
+   this.transformDataStep = new aws_stepfunctions.Pass(
      this,
      "TransformIncomingDataStep",
      {
        parameters: {
          payload: aws_stepfunctions.JsonPath.stringAt(
            "States.Format('{} {}', $.firstName, $.lastName)"
          )
        }
      }
    );

    // You can imagine the definition being a bit more complex.
-   const stepFunctionDefinition = transformDataStep;
+   const stepFunctionDefinition = this.transformDataStep;

    const stepFunction = new aws_stepfunctions.StateMachine(
      this,
      "StepFunction",
      {
        definition: stepFunctionDefinition
      }
    );
  }
}
```

With the `transformDataStep` made public, the test body would look as follows.

```ts
// pass-states-integration.test.ts

// Test setup ...

test("Handles the failure of the BackgroundCheck step", async () => {
  const stack = new cdk.Stack();
  const construct = new PassStatesIntegration(stack, "PassStatesIntegration");

  const transformDataStepDefinition = construct.transformDataStep.toStateJson();
  const stepFunctionDefinition = JSON.stringify({
    StartAt: "TransformDataStep",
    States: {
      TransformDataStep: {
        ...transformDataStepDefinition,
        End: true
      }
    }
  });

  const sfnLocalClient = new SFNClient({
    endpoint: `http://localhost:${container?.getMappedPort(8083)}`
  });

  const createLocalSFNResult = await sfnLocalClient.send(
    new CreateStateMachineCommand({
      definition: stepFunctionDefinition,
      name: "PassStates",
      roleArn: "arn:aws:iam::012345678901:role/DummyRole"
    })
  );

  const startLocalSFNExecutionResult = await sfnLocalClient.send(
    new StartExecutionCommand({
      stateMachineArn: createLocalSFNResult.stateMachineArn,
      input: JSON.stringify({
        firstName: "John",
        lastName: "Doe"
      })
    })
  );

  await waitFor(async () => {
    const getExecutionHistoryResult = await sfnLocalClient.send(
      new GetExecutionHistoryCommand({
        executionArn: startLocalSFNExecutionResult.executionArn
      })
    );

    const successState = getExecutionHistoryResult.events?.find(
      event => event.type == "ExecutionSucceeded"
    );

    expect(successState?.executionSucceededEventDetails?.output).toEqual(
      JSON.stringify({ payload: "John Doe" })
    );
  });
}, 20_000);
```

The ASL for the `TransformDataStep` is extracted via the `toStateJson` method. The rest of the test is similar to how we did it previously. The only difference is how we make the assertion.

This testing method is analogous to the one described by the _By decomposing AWS Step Function Tasks_ section.

## Closing words

I hope you find this blog post helpful regarding AWS Step Functions testing.

Consider following me on Twitter for more serverless content - [@wm_matuszewski](https://twitter.com/wm_matuszewski).

Thank you for your valuable time.