To mock the Lambda functions during interacting with the StepFunctions Local, a solution is creating a fake Lambda HTTP service in a Python thread initiated at the testing setup and making this service able to parse the HTTP request URL to determine which function to invoke.
I've implemented this concept as a pytest fixture: https://github.com/chehsunliu/pytest-stepfunctions.
Usage
Suppose there is a state machine which simply collects all the EMR cluster Ids and we want to test it locally.
State Machine Definition
{
"StartAt": "ListIds",
"States": {
"ListIds": {
"Type": "Task",
"Resource": "${ListIdsLambdaArn}",
"ResultPath": "$.cluster_ids",
"End": true
}
}
}
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"]]
Test Code
tests/test_foo.py
import json
import time
from string import Template
import boto3
from botocore.stub import Stubber
def test_bar(aws_stepfunctions_endpoint_url):
# Create the definition string.
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)
# Create the state machine resource.
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"]
# Mock the Lambda code.
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"}]}
)
# Start and wait until the execution finishes.
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)
# Validate the results.
stubber.assert_no_pending_responses()
assert "SUCCEEDED" == response["status"]
assert ["j-00001", "j-00002"] == json.loads(response["output"])["cluster_ids"]
Running the test
Install the dependencies:
$ pip install boto3 pytest pytest-stepfunctions pytest-mock
Download the StepFunctions Local JAR here and execute it:
$ java -jar /path/to/StepFunctionsLocal.jar \
--lambda-endpoint http://localhost:13000 \
--step-functions-endpoint http://localhost:8083 \
--wait-time-scale 0
Run the test:
$ 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
The test can also be performed in Docker Compose, which is much easier to use and maintain. You can check the README in my repo. Hope this fixture could help people who found this article.
