Thursday, March 26, 2020

Mock boto3 services not handled by moto

We have a pattern to mock AWS services in our python lambdas. First, in our lamba code, we initialize boto3 clients with properties:

import os
import boto3

@property
def REGION_NAME():
    return os.getenv("REGION", "eu-west-3")

@property
def SQS():
    return boto3.client(service_name="sqs", region_name=REGION_NAME.fget())

def lambda_handler(event, context):
    sqs_client = SQS.fget()
As you can see, the region is also a property. The reason for it is that all clients in moto are declared on the us-east-1 region.
Our unit tests are all in a test sub-folder of our lambda. To write the tests, we usually follow this pattern:
from moto import mock_sqs
from pytest import fixture

from ..mylambda import (
    lambda_handler,
    SQS
)

@fixture(autouse=True)
def prepare_test_env(monkeypatch):
    monkeypatch.setenv("REGION", "us-east-1")   

@mock_sqs
def test_mylambda():
    SQS.fget().create_queue(QueueName='MyQueue')
    
    #do the test...
We import from the lambda in the parent folder the methods to test, as well as the properties. For this to work, we have to create an empty __init__.py file in that directory. We then patch the environment variables, including the region that needs setting to us-east-1. In our test function, we have to mock the boto3 client, using the appropriate moto annotation. Then we write our test code.

In some cases, the boto3 client has no mock in moto. That is our case for Step Functions for instance (I know it is in preparation, but not yet ready as the time of the writing). For those cases, we use the following pattern:

import boto3

@property
def SFN():
    return boto3.client("stepfunctions")
In our lambda, nothing changes. We are still using a property. However, in the unit test, we have to use a patch:

from unittest.mock import (
    PropertyMock,
    patch
)
from ..mylambda import lambda_handler

def test_lambda():
    with patch('mylambda.mylambda.SFN', new_callable=PropertyMock) as mock_stepfunctions:
        lambda_handler(event)

        mock_stepfunctions.fget().start_execution.assert_called_with(
            stateMachineArn="MyMachine", 
            name="State-Machine-0", 
            input="{}"
        )
    
We patch the property SFN with a PropertyMock. By giving it a name, we can then use the mock to assert that it was called with the correct parameters.

No comments:

Post a Comment