Test your Step Functions workflows locally with pytest

Before you deploy a Step Functions workflow, it is hard to check what happens when a service fails: which errors get retried, which Catch handles them, and what the next Task receives. AWS documents Step Functions Local as unsupported. The recommended alternative, the TestState API, supports mocks but tests one state per call. To test a complete path, you have to chain those calls yourself.

sfnx includes a local runner for testing an entire workflow with pytest. You provide mock service responses, the runner follows the definition's Retry and Catch rules, and your assertions check the calls and the result. It works with existing JSONata definitions, so there's no need to rewrite the workflow or add states to it.

sfnx also has a compiler if you'd prefer to maintain the workflow itself in Python. I'll cover that option in the second half.

A test can run the ASL you already have

The runner reads the same JSON that Step Functions runs. A workflow is a JSON document in the Amazon States Language (ASL), whether you write it by hand or build it in Workflow Studio, the console's visual editor. In JSONata mode, fields such as a Task's Arguments can hold expressions inside {% ... %}.

For this example, I'll use AWS's coding-agent sample. It looks up a clinical term with Lambda, calls an AI agent if there's no direct match, routes the answer by confidence, and writes back the result. The agent runs in an Amazon Bedrock AgentCore harness, which Step Functions can invoke directly as a Task.

Each Task has a Retry policy for transient errors. If the lookup, agent call or initial write-back fails with an error handled by its Catch, the workflow moves to a fallback Task that marks the term as open and records the reason. If the fallback also fails after its retries, the execution fails.

The test below asks one question: if the agent is throttled on every attempt, how many times is it called, and what does the fallback Task receive?

Set up a project

You need uv and Python 3.11 or later. In a new directory:

uv init
uv add --dev "sfnx[testing]==2.1.0" pytest

These examples were tested with sfnx v2.1.0. Omit ==2.1.0 to install the latest version instead.

Save the JSON form of the sample's definition from sfnx v2.1.0 as coding-workflow.asl.json in the project root. It is AWS's definition, converted from the sample's YAML, with nothing added.

Write the test

One mock function handles the lookup, agent call and write-back. Save this test as test_coding_workflow.py in the project root:

import json
from pathlib import Path

from sfnx.testing import Call, Failure, run

DEFINITION = json.loads(Path("coding-workflow.asl.json").read_text())
TERM = {
    "record_id": "r1",
    "verbatim": "migrane",
    "encoding_dictionary": "MedDRA",
    "encoding_dictionary_version": "v27.0",
    "source_study": "ONCO-2024-01",
}
HARNESS = "arn:aws:states:::bedrockagentcore:invokeHarness"


def throttled_agent(call: Call) -> object:
    """No direct match, the agent is throttled, and the write-back succeeds."""
    if call.resource == HARNESS:
        raise Failure("BedrockAgentCore.ThrottlingException", "slow down")
    assert isinstance(call.arguments, dict)
    if call.arguments["FunctionName"] == "${CheckDirectFunctionArn}":
        return {"Payload": {"blocked": False, "matched": False}}
    return {"Payload": call.arguments["Payload"]}


def test_a_throttled_agent_is_retried_then_the_term_is_left_open():
    execution = run(DEFINITION, TERM, throttled_agent)
    harness_calls = [c for c in execution.calls if c.resource == HARNESS]
    assert len(harness_calls) == 4  # the first attempt and 3 retries
    last = execution.calls[-1].arguments
    assert last["Payload"]["target_status"] == "open"
    assert last["Payload"]["failure_reason"] == "BedrockAgentCore.ThrottlingException"
    assert execution.output == {
        "record_id": "r1",
        "target_status": "open",
        "failure_reason": "BedrockAgentCore.ThrottlingException",
        "rationale": None,
    }

Run it from the project root:

uv run pytest test_coding_workflow.py

What the runner did

This test runs locally without AWS credentials or calls to AWS services. run follows the definition state by state. At each Task, it passes the Resource and evaluated Arguments to throttled_agent and uses the return value as the service response.

When the mock raises Failure, the runner applies the definition's error handling. Here, the agent gets one initial attempt and three retries before the Catch routes execution to the fallback Task.

The assertions check the four agent attempts, the fallback Task's arguments, and the final output. If the execution fails, accessing execution.output raises the error, which also fails the test.

${...} placeholders are left unchanged. The mock receives ${CheckDirectFunctionArn} as written, so you can test the definition before substituting the real ARNs.

The runner skips delays: retries don't wait between attempts, and Wait states return immediately. Testing a retry policy doesn't mean waiting through its backoff intervals.

Testing a definition from Workflow Studio

To test your own JSONata workflow, export its definition from the console as JSON. Use each Task's Resource and Arguments to decide what your mock should return, or raise Failure to test an error path. Then add assertions for the calls and output you expect.

I also tested a small JSONata workflow built in Workflow Studio and exported without edits. It had a Lambda call with the retry policy the editor adds, a Choice, Succeed and Fail states, and a Catch.

It followed the expected path in all four cases I tried: success, the Choice's other branch, a throttled call caught after retries, and an error caught without retries. I haven't tested every state type or integration available in the editor.

The runner checks for unsupported states and fields before execution. It accepts JSONata mode only; a JSONPath state or an unsupported field produces an error identifying the state. See the testing guide for the supported states and fields.

A local run checks control flow against your mocks, not the services. It does not reproduce elapsed time, concurrency, real service behavior, or every detail of JSONata on AWS. The testing guide lists where a local run differs.

If you want to maintain the workflow in Python, sfnx can compile it

The compiler is optional; nothing above used it. You write the workflow's intent in Python, sfnx compiles it to readable ASL in JSONata mode, and you deploy the generated definition with your usual tools.

Here's how the same coding-agent workflow looks in Python. You can find the full source and the generated definition in the repository.

Define a shared retry policy once

The original definition repeats the same Lambda retry policy on four Tasks. In Python, it is one module-level constant (excerpt; the full constant has a second retrier for timeouts and database errors):

LAMBDA_RETRY = [
    {
        "ErrorEquals": [Lambda.TooManyRequestsException, Lambda.ServiceException],
        "IntervalSeconds": 2,
        "MaxAttempts": 3,
        "BackoffRate": 2.0,
    },
    ...
]

Each Lambda Task uses the same constant:

direct = task(
    "arn:aws:states:::lambda:invoke",
    {
        "FunctionName": "${CheckDirectFunctionArn}",
        "Payload": {
            "verbatim": verbatim,
            "encoding_dictionary": dictionary,
            "encoding_dictionary_version": version,
        },
    },
    retry=LAMBDA_RETRY,
)["Payload"]

To change the backoff, you edit the constant once. The compiler includes the full policy in each Task that uses it.

Write branches with if/elif

The original uses a Choice state to select the next Task based on the agent's score:

"ScoreThreshold": {
  "Type": "Choice",
  "Choices": [
    { "Condition": "{% $agentScore >= 0.9 %}", "Next": "WriteBack" },
    { "Condition": "{% $agentScore >= 0.7 %}", "Next": "MarkForReview" }
  ],
  "Default": "MarkOpen"
}

In Python, the same decision is an if / elif / else:

if score >= 0.9:
    status = "autocoded"
elif score >= 0.7:
    status = "approval_required"
else:
    status = "open"

This compiles to a single Choice state. Each matching rule sets the status variable; the default path sets it to open:

"if_3": {
  "Type": "Choice",
  "Choices": [
    {
      "Condition": "{% $score >= 0.9 %}",
      "Assign": { "status": "autocoded" },
      "Next": "if_4"
    },
    {
      "Condition": "{% $score >= 0.7 %}",
      "Assign": { "status": "approval_required" },
      "Next": "if_4"
    }
  ],
  "Assign": { "status": "open" },
  "Default": "if_4"
}

The Python version's graph has a different shape from the original: it uses one write-back Task for the accepted and review cases, passing status as target_status.

Handle errors with try/except: only Task, Parallel and Map states catch

There is an important difference from Python's try / except: the compiler adds Catch clauses to Task, Parallel and Map states in the try block, but Pass states cannot catch errors.

The example therefore parses the agent's reply in the same statement as the task() call. If parsing fails, the Task's Catch handles the error. Moving the parsing to a separate statement would put it in a Pass state, where a parsing error would fail the entire execution. The language reference covers the details.

Stay within the supported Python subset

The compiler supports a subset of Python, so arbitrary library calls won't compile. Unsupported constructs produce an error with a suggested alternative. You can also use jsonata() to write an expression directly in JSONata.

Deployment stays with your existing tools. Resources, IAM and substitutions go through CDK, SAM, CloudFormation or the AWS CLI; the deployment guide has examples.

The Python version behaves like the original in 12 mocked scenarios

To check the Python version against the original, this test runs both definitions through 12 scenarios with the same mocks. It uses AWS's definition from commit ab76ef3 (MIT-0) and compares every call's resource and arguments, along with the final output or error.

The scenarios cover direct matches, blocked terms, each confidence band, replies without a score, a rationale or JSON, and each service failing. To reproduce them:

git clone https://github.com/iwamot/sfnx.git
cd sfnx
git checkout v2.1.0
uv run pytest tests/test_samples.py

Both definitions produce the same calls and results for these scenarios in the local runner. That doesn't establish that they behave identically on AWS with real Lambda and AgentCore responses.

The compiler's trade-off: it can add states

In this example, the compiled workflow has more states than the original. This affects workflows you compile with sfnx; testing existing ASL with the runner adds no states. The compiler emits Choice states for branches and separate states for assignments it cannot combine with another state.

Some assignments are combined with an existing state: a Wait can hold the assignments after it, a Task can assign its result, and a Choice rule can set a variable. The compiler keeps states separate when combining them could change the behavior.

Across the 12 scenarios, the Python version enters 1.5 to 2.3 times as many states on the same path.

These are local execution counts, not billing estimates. The runner records each state entry but excludes retry attempts from this count. Standard workflows are billed per state transition, including retries.

What you pay depends on the path each execution takes and how many executions you run. For a Standard workflow that runs often, look at the generated graph before you adopt the compiler.

For a smaller example, see the Python version of the console's Hello World template. It shows how if, wait(), parallel() and raise compile. The generated definition has 11 states compared with the template's 9; a successful run enters 9 states compared with 8.

Start with one failure path

Try adding a test for one failure path in a JSONata workflow you already have. Export the definition, mock its Tasks, and check that the workflow handles the error and passes the expected data to the next Task.

If you'd also like to maintain the workflow in Python, start with a small example and inspect the generated ASL before deciding whether to use the compiler.

sfnx is on GitHub. If you try it on a real workflow, I'd like to hear which paths you tested, and where the runner fell short.