Unit testing AWS services in Python

Consider the following piece of code:

import boto3
Table = boto3.resource('dynamodb').Table('foo')
def get_user(user_id):
ddb_response = Table.get_item(Key={'id': user_id})
return ddb_response.get('Item')
view raw models.py hosted with ❤ by GitHub

It’s a contrived example that just reads an item of data from a DynamoDB table. How would you write a unit test for the get_user function?

My favourite way to do so is to combine pytest fixtures and botocore’s Stubber:

from botocore import Stubber, ANY
import pytest
import models
@pytest.fixture(scope="function")
def ddb_stubber():
ddb_stubber = Stubber(models.Table.meta.client)
ddb_stubber.activate()
yield ddb_stubber
ddb_stubber.deactivate()
def test_user_exists(ddb_stubber):
user_id = 'user123'
get_item_params = {'TableName': ANY,
'Key': {'id': user_id}}
get_item_response = {'Item': {'id': {'S': user_id},
'name': {'S': 'Spam'}}}
ddb_stubber.add_response('get_item', get_item_response, get_item_params)
result = main.get_user(user_id)
assert result.get('id') == user_id
ddb_stubber.assert_no_pending_responses()
def test_user_missing(ddb_stubber):
user_id = 'user123'
get_item_params = {'TableName': ANY,
'Key': {'id': user_id}}
get_item_response = {}
ddb_stubber.add_response('get_item', get_item_response, get_item_params)
result = main.get_user(user_id)
assert result is None
ddb_stubber.assert_no_pending_responses()
view raw test_models.py hosted with ❤ by GitHub

There’s couple of things to note here.

First, I’m using the wonderful scope functionality of pytest fixtures. This allows me to create a new fixture per every test function execution. It is necessary for Stubber to work correctly.

The Stubber needs to be created with the correct client. Since I’m using a DynamoDB Table instance in models.py, I have to access its client when creating the Stubber instance.

Notice also the “verbose” get_item_response structure in the first test. That’s because of how the DynamoDB client interacts with DynamoDB API (needless to say, this is DynamoDB specific). The Table is a layer of abstraction on top of this, it converts between DynamoDB types and Python types. However it still uses the client underneath, so it expects this structure nevertheless.

Finally, it’s good practice to call assert_no_pending_response to make sure the tested code actually did make the call to an AWS service.

I really like this combination of pytest and Stubber. It’s a great match for writing correct and compact tests.

Share your thoughts

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Google photo

You are commenting using your Google account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s