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.
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:
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