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') |
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() |
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.
