Mocking multiple boto3 services, some without moto implementation

Question

I am trying to unit test the logic in a AWS Lambda function using mocking. The Lambda finishes it's execution by sending push notifications via AWS Pinpoint. The Lambda also uses AWS SSM Parameter Store. I have been mocking in other Lambdas, with multiple boto3 objects, with moto https://github.com/spulec/moto but there is no Pinpoint implementation in moto at present.

I found a solution in https://stackoverflow.com/a/55527212/839338 which I needed to modify to get it to work. The question it was answering was not about my exact scenario but the answer pointed me to a solution. So I am posting here to document my changes to the solution which I modified and to ask if there is a more elegant way of doing this. I have looked at botocore.stub.Stubber but can't see a way it is better, but I'm willing to be proved wrong.

My code so far:

test.py

import unittest
from unittest.mock import MagicMock, patch
import boto3
from moto import mock_ssm
import my_module


def mock_boto3_client(*args, **kwargs):
    if args[0] == 'ssm':
        # Use moto.
        mock_client = boto3.client(*args, **kwargs)
    else:
        mock_client = boto3.client(*args, **kwargs)
        if args[0] == 'pinpoint':
            # Use MagicMock.
            mock_client.create_segment = MagicMock(
                return_value={'SegmentResponse': {'Id': 'Mock SegmentID'}}
            )
            mock_client.create_campaign = MagicMock(
                return_value={'response': 'Mock Response'}
            )
    return mock_client


class TestMyModule(unittest.TestCase):
    @patch('my_module.boto3')
    @mock_ssm
    def test_my_module(self, mock_boto3):
        mock_boto3.client = mock_boto3_client
        conn = mock_boto3.client('ssm', region_name='eu-west-2')
        conn.put_parameter(
            Name='/my/test',
            Value="0123456789",
            Type='String',
            Tier='Standard'
        )
        response = my_module.handler()
        self.assertEqual(
            ('0123456789', 'Mock SegmentID', {'response': 'Mock Response'}), 
            response
        )

my_module.py

import boto3
import json


def get_parameter():
    ssm = boto3.client('ssm', region_name='eu-west-2')
    parameter = ssm.get_parameter(Name='/my/test')
    return parameter['Parameter']['Value']


def create_segment(client, message_id, push_tags, application_id):
    response = client.create_segment(
        ApplicationId=application_id,
        WriteSegmentRequest={
            'Dimensions': {
                'Attributes': {
                    'pushTags': {
                        'AttributeType': 'INCLUSIVE',
                        'Values': push_tags
                    }
                }
            },
            'Name': f'Segment {message_id}'
        }
    )
    return response['SegmentResponse']['Id']


def create_campaign(client, message_id, segment_id, application_id):
    message_payload_apns = json.dumps({
        "aps": {
            "alert": 'My Alert'
        },
        "messageId": message_id,
    })

    response = client.create_campaign(
        ApplicationId=application_id,
        WriteCampaignRequest={
            'Description': f'Test campaign - message {message_id} issued',
            'MessageConfiguration': {
                'APNSMessage': {
                    'Action': 'OPEN_APP',
                    'RawContent': message_payload_apns
                }
            },
            'Name': f'{message_id} issued',
            'Schedule': {
                'StartTime': 'IMMEDIATE'
            },
            'SegmentId': segment_id
        }
    )
    return response


def handler():
    application_id = get_parameter()
    client = boto3.client('pinpoint', region_name='eu-west-1')
    segment_id = create_segment(client, 12345, [1, 2], application_id)
    response = create_campaign(client, 12345, segment_id, application_id)
    return application_id, segment_id, response

In particular I would like to know how to implement mock_boto3_client() better and more elegantly to handle in a more generic way.

Bert Blommers · Accepted Answer

It's relatively easy to use the moto framework for any new services. This allows you to focus on the required behaviour, and moto takes care of the scaffolding.

There are two steps required to register an additional service in the Moto-framework:

Ensure that moto mocks the actual HTTP requests to https://pinpoint.aws.amazon.com
Create a Responses class that acts on the requests for https://pinpoint.aws.amazon.com

Mocking the actual HTTP requests can be done by extending the BaseBackend-class from moto. Note the urls, and the fact that all requests to this url will be mocked by the PinPointResponse-class.

pinpoint_mock/models.py:

import re

from boto3 import Session

from moto.core import BaseBackend
from moto.sts.models import ACCOUNT_ID



class PinPointBackend(BaseBackend):

    def __init__(self, region_name):
        self.region_name = region_name

    @property
    def url_paths(self):
        return {"{0}/$": PinPointResponse.dispatch}

    @property
    def url_bases(self):
        return ["https?://pinpoint.(.+).amazonaws.com"]

    def create_app(self, name):
        # Store the app in memory, to retrieve later
        pass


pinpoint_backends = {}
for region in Session().get_available_regions("pinpoint"):
    pinpoint_backends[region] = PinPointBackend(region)
for region in Session().get_available_regions(
    "pinpoint", partition_name="aws-us-gov"
):
    pinpoint_backends[region] = PinPointBackend(region)
for region in Session().get_available_regions("pinpoint", partition_name="aws-cn"):
    pinpoint_backends[region] = PinPointBackend(region)

The Response-class needs to extend the BaseResponse-class from moto, and needs to duplicate the method-names that you're trying to mock.
pinpoint/responses.py

from __future__ import unicode_literals

import json

from moto.core.responses import BaseResponse
from moto.core.utils import amzn_request_id
from .models import pinpoint_backends


class PinPointResponse(BaseResponse):
    @property
    def pinpoint_backend(self):
        return pinpoint_backends[self.region]

    @amzn_request_id
    def create_app(self):
        name = self._get_param("name")
        pinpoint_backend.create_app(name)
        return 200, {}, {}

Now all that's left is creating a decorator:

from __future__ import unicode_literals
from .models import stepfunction_backends
from ..core.models import base_decorator

pinpoint_backend = pinpoint_backends["us-east-1"]
mock_pinpoint = base_decorator(pinpoint_backends)

@mock_pinpoint
def test():
    client = boto3.client('pinpoint')
    client.create_app(Name='testapp')

The code was taken from the StepFunctions-module, which is probably one of the simpler modules, and easiest to adapt to your needs: https://github.com/spulec/moto/tree/master/moto/stepfunctions

Mocking multiple boto3 services, some without moto implementation

Answers (2)

Related Questions