Skip to content

Repository files navigation

AWS Step Functions Toolkit

Run a Step Functions state machine end-to-end on your laptop — and choose, per step, how each one runs.

Documentation Status

The motivation to create this package was the need to test a Step Functions pipeline built almost entirely from batch:submitJob.sync steps. AWS's test_state API does not currently support .sync integrations or .waitForTaskToken, and deploying and running the state machine on AWS to validate each change made iterating impractical.

AWS's docs suggest chaining test_state calls — feeding each state's output and next state into the next call — to exercise a complete execution path, but leave you to build that driver loop yourself. This toolkit packages that logic as a framework so you don't have to: it walks your real ASL definition state by state, delegates the engine logic (Path/Parameters/ResultSelector/Output and next-state resolution) to test_state, and handles the parts that aren't a linear chain — Map and Parallel fan-out, any-depth recursion, and nested subflows. For the steps test_state cannot run (.sync, .waitForTaskToken), you plug in a strategy — typically by building and running that step's container locally with Docker. The result: run an entire pipeline locally and swap any step between a mock, your own function, a local container, or a real AWS call.

Annotated Step Functions workflow showing how the toolkit runs each state type: test_state for engine logic, a strategy (e.g. local Docker) for .sync steps, and recursion for nested subflows

The real console graph of the docker-batch example, annotated to show how the toolkit handles each state.

Install

pip install aws-stepfunctions-toolkit

Quick start

To run the bundled quickstart example (no Docker — its task steps are mocked), head to examples/quickstart/ — it has step-by-step instructions. You'll need AWS credentials + a region and an IAM role allowed to call test_state (see the Setup guide).

What your code looks like

The snippet below shows the shape of a run — it is not runnable on its own (you supply your own state machine); the quickstart folder above is the runnable version.

import json
from aws_stepfunctions_toolkit import WorkflowRunner, StaticMockResponseStrategy, CallableStrategy

# Bring your own: your Step Functions state machine, exported as Amazon States Language (ASL)
# JSON. (In the AWS console: the state machine's "Definition" tab. This is the workflow you
# want to test — the toolkit does not generate it for you.)
definition = json.loads(open("state_machine.asl.json").read())

# For each step you don't want to run for real, say how to produce its result. Keys are the
# state names from your definition.
mock_mapping = {
    # A mocked lambda:invoke result puts the function's return under "Payload" as a JSON string.
    "Enrich": CallableStrategy(lambda data: {"Payload": json.dumps({"enriched": True})}),          # your own function
    "Notify": StaticMockResponseStrategy(json.dumps({"Payload": json.dumps({"status": "sent"})})),  # a fixed payload
}

runner = WorkflowRunner(
    # Each state machine carries its own execution role (the role its states run under);
    # this is the only value you must set — needs AWS creds + a region in your environment.
    asl_registry={
        "main": {**definition, "ROLE_ARN": "arn:aws:iam::<account>:role/<role-with-test-state-perms>"},
    },
    mock_mapping=mock_mapping,
)

final_output = runner.start(initial_input={"order_id": 123})
print(final_output)

asl_registry maps state-machine names to their ASL definitions: put the workflow you're testing under the key "main" (the entry point). If your machine starts nested state machines (via startExecution.sync:2), register each one too, keyed by the name of the state that starts it — see Control flow.

JSONata throughout. The examples (and the toolkit's mock-result handling) use the JSONata query language. Set "QueryLanguage": "JSONata" on each state (not just at the top level): the runner tests one state at a time, so a state must declare its own query language.

States without an entry in mock_mapping are handled automatically (test_state for ordinary states; built-in recursion for Map / Parallel / nested state machines). To run a step in a real local container instead of mocking it, use DockerBatchStrategy (for Batch/.sync steps) or DockerLambdaStrategy (to run a lambda:invoke step's real image via the Runtime Interface Emulator).

Features

Each links to docs with a code snippet, or to a runnable example.

How this compares to existing tools

Two options exist for running Step Functions locally: AWS Step Functions Local and LocalStack, a general-purpose local-AWS emulator.

This toolkit takes a deliberately narrower approach — complementary to, not a replacement for, a full local-cloud emulator:

  • A library, not an environment. You run your real ASL directly against AWS's real test_state engine — nothing to provision, no separate stack to stand up or keep in sync.
  • Runs your real step. Each step runs locally via a strategy — your actual container via Docker, your own Python function, a subprocess, a static mock, or a real AWS call.
  • Partial runs. Because it drives the walk itself, it can start and stop at specific states, so you can re-run just part of a long pipeline.
  • Real callbacks. .waitForTaskToken steps — e.g. Batch jobs that return their result via the task token — run for real, not only as a static mock. (LocalStack's Batch integration does not support .waitForTaskToken.)

On the existing options: AWS notes that testing state machines with Step Functions Local is unsupported and points to the test_state API instead. LocalStack looks great, and it may be the better fit if you want a full local AWS environment. We wanted the opposite for this use case: a lightweight solution we have full control over and can customize — purpose-built for fast, low-setup iteration on a Step Functions pipeline, not a general environment to run.

Documentation

Full documentation is hosted at aws-stepfunctions-toolkit.readthedocs.io.

Page What's in it
Setup AWS credentials, the test_state IAM role, and Docker.
How it works Motivation, the test_state gap, and the run loop.
Usage walkthrough A complete, end-to-end worked example.
Selecting how each step runs The strategy catalog + image sources.
Control flow Subflows, Map/Parallel, recursion, start/end control.
Container-side handler BatchJobInterface contract.
CLI & history sfn-toolkit generate-mock, ExecutionHistory.

Requirements

  • Python 3.13+
  • An AWS role/credentials allowed to call test_state (the toolkit calls the real API for the engine logic).
  • Docker, only if you use DockerBatchStrategy to run steps in containers.

New to this? The Setup guide covers AWS credentials, creating the test_state IAM role, and Docker, step by step.

Examples

Each folder under examples/ is self-contained (its own ASL, runnable script, and README) — see the examples index. Start with:

  • examples/quickstart/ — copy-and-run starter; set ROLE_ARN and python run.py. No Docker (task steps are mocked).
  • examples/local-subprocess/ — run a step's code as a local subprocess (LocalExecutionStrategy). No Docker.
  • examples/docker-batch/ — run steps in real local containers (DockerfileImage + BakeImage) plus a nested machine. make run-example runs it (needs Docker, AWS creds, and a ROLE_ARN env var).
  • examples/docker-lambda/ — run a lambda:invoke step's real Lambda container image via the Runtime Interface Emulator (DockerLambdaStrategy). Needs Docker.

License

See LICENSE.

About

Package provides convenient tools for AWS Step Functions testing.

Topics

Resources

Stars

1 star

Watchers

1 watching

Forks

Releases

Packages

Contributors

Languages