Skip to main content

A small description

Project description

pytest-stepfunctions

GitHub Actions/CI codecov pypi-version pypi-downloads pypi-pyversions

A pytest fixture that makes you able to mock Lambda code during AWS StepFunctions local testing.

Table of Contents

Overview

AWS provides local Step Functions as a JAR and a Docker image for the quick testing without deployment. They described how to perform such task in this article as well. I got excited at the very beginning, but soon ended up frustrated for still being unable to mock Lambda functions' external dependencies. Then I thought: what if initiate a Python thread with a fake Lambda service and use this fake service to execute Lambda functions? Fortunately, It works!

Installing

Use pip to install:

$ pip install pytest-stepfunctions

Getting Started

Suppose there is a state machine that simply collects all the EMR cluster unique identifiers. Here is the state machine definition:

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

and the 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"]]

Creating a State Machine

In the test file tests/test_foo.py, create a Step Functions client with endpoint URL pointed to our Step Functions service, and use this client to create a state machine resource by using the definition above

from string import Template

import boto3


def test_bar(aws_stepfunctions_endpoint_url):
    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)

    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"]

Note the internal fake Lambda service in pytest-stepfunctions will parse Lambda ARNs to recognize what to call.

Mocking the EMR Client in the Lambda Code

Here uses the pytest-mock fixture to temporarily patch the boto3 module inside the Lambda code. botocore.stub.Stubber is also applied to make sure the mock request parameters and response content are all valid:

from botocore.stub import Stubber


def test_bar(aws_stepfunctions_endpoint_url, mocker):
    ...

    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"}]}
    )

Starting Execution and Validating Results

Start and wait until the execution status is not RUNNING:

import json
import time


def test_bar(aws_stepfunctions_endpoint_url, mocker):
    ...

    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)

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

Running the Test with the Step Functions JAR

The JAR is available here. Download and execute it first:

$ java -jar /path/to/StepFunctionsLocal.jar \
    --lambda-endpoint http://localhost:13000 \
    --step-functions-endpoint http://localhost:8083 \
    --wait-time-scale 0
Step Functions Local
Version: 1.4.0
Build: 2019-09-18
2020-07-06 18:40:28.284: Configure [Lambda Endpoint] to [http://localhost:13000]
2020-07-06 18:40:28.285: Configure [Step Functions Endpoint] to [http://localhost:8083]
2020-07-06 18:40:28.323: Loaded credentials from profile: default
2020-07-06 18:40:28.324: Starting server on port 8083 with account 123456789012, region us-east-1

Then run the test with the following command:

$ 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
=============================== test session starts ================================
platform linux -- Python 3.7.3, pytest-5.4.3, py-1.9.0, pluggy-0.13.1 -- /tmp/gg/venv/bin/python
cachedir: .pytest_cache
rootdir: /tmp/gg
plugins: mock-3.1.1, stepfunctions-0.1a2
collected 1 item

tests/test_foo.py::test_bar PASSED                                           [100%]

================================ 1 passed in 1.01s =================================

Running the Test with the Step Functions Docker Image

I personally recommend this way as it is much easier to reproduce the testing environment.

This is the Dockerfile

FROM python:3.7

WORKDIR /app

COPY ./my ./my
COPY ./tests ./tests
RUN pip install pytest pytest-stepfunctions pytest-mock boto3

and the docker-compose.yml for Docker Compose:

version: "3.2"

services:
  tester:
    build:
      context: .
      dockerfile: ./Dockerfile
    environment:
      AWS_DEFAULT_REGION: us-east-1
      AWS_ACCESS_KEY_ID: xxx
      AWS_SECRET_ACCESS_KEY: xxx
    command: >
      bash -c "python -m pytest -v
      --pytest-stepfunctions-endpoint-url=http://sfn-endpoint:8083
      --pytest-stepfunctions-lambda-address=0.0.0.0
      --pytest-stepfunctions-lambda-port=13000
      ./tests"

  sfn-endpoint:
    image: amazon/aws-stepfunctions-local:1.5.1
    environment:
      AWS_DEFAULT_REGION: us-east-1
      AWS_ACCESS_KEY_ID: xxx
      AWS_SECRET_ACCESS_KEY: xxx
      WAIT_TIME_SCALE: 0
      STEP_FUNCTIONS_ENDPOINT: http://sfn-endpoint:8083
      LAMBDA_ENDPOINT: http://tester:13000

Then run the following command to run the test:

$ docker-compose up --build --exit-code-from tester

Contributing

Here are the tools required:

  • Docker and Docker Compose
  • GNU Make
  • pip-tools

Set up a virtual environment for your IDE:

$ python -m venv venv
$ source ./venv/bin/activate
(venv) $ pip-sync ./dev-requirements.txt ./requirements.txt

Run the tests and linters:

$ make lint test

Known Issues

  1. Nested workflows are very slow: if a state machine contains lots of nested state machines, the execution will be extremely slow even WAIT_TIME_SCALE is set to 0. This is a known performance issue in the official JAR.
  2. AWS Service integrations other than Lambda are not supported yet. Services like EMR even have no endpoint option in the official JAR. A possible workaround for some cases is calling them by invoking Lambda functions.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

pytest-stepfunctions-0.1.1.tar.gz (25.9 kB view details)

Uploaded Source

Built Distribution

pytest_stepfunctions-0.1.1-py3-none-any.whl (8.1 kB view details)

Uploaded Python 3

File details

Details for the file pytest-stepfunctions-0.1.1.tar.gz.

File metadata

  • Download URL: pytest-stepfunctions-0.1.1.tar.gz
  • Upload date:
  • Size: 25.9 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/3.4.1 importlib_metadata/4.0.1 pkginfo/1.7.0 requests/2.25.1 requests-toolbelt/0.9.1 tqdm/4.60.0 CPython/3.9.5

File hashes

Hashes for pytest-stepfunctions-0.1.1.tar.gz
Algorithm Hash digest
SHA256 b3d1a2df7cef142f235a4be44d8c69ee60175db6de9385d5a9a2d2242af7b914
MD5 e7910bfd2cb5ba2fecd0707d8e92f008
BLAKE2b-256 3025607910bfb5c612a84f9b7ec08cc6de25a480aad3785c7a23bbb956a5dec0

See more details on using hashes here.

File details

Details for the file pytest_stepfunctions-0.1.1-py3-none-any.whl.

File metadata

  • Download URL: pytest_stepfunctions-0.1.1-py3-none-any.whl
  • Upload date:
  • Size: 8.1 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/3.4.1 importlib_metadata/4.0.1 pkginfo/1.7.0 requests/2.25.1 requests-toolbelt/0.9.1 tqdm/4.60.0 CPython/3.9.5

File hashes

Hashes for pytest_stepfunctions-0.1.1-py3-none-any.whl
Algorithm Hash digest
SHA256 76bd42b718bf345ef9501b3a024dbad3fd7d3d7ebf0a7a6a5e9d4f3cb6ad3da7
MD5 5add31ba73e63cfcfb9e0659322e0ecf
BLAKE2b-256 1769ca49641a6aeca2eda6e024d09b46931ab44f6ebd58200ea594a0db35af00

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page