Unit Testing Python Code that Uses the AWS SDK [how-to]

Intro to unit testing Python code that incorporates the AWS SDK

Unit Testing in Python

Creating unit tests in Python is an excellent way to not only regression test code but also help with development. There are a number of testing frameworks that include unittest from the core library and others available outside the core library such as pytest. Pytest is a popular and easy-to-use testing framework and I’ve chosen to use it in this post. It doesn’t require the extension of any classes so you can immediately start writing tests instead of reading documentation.

If you’re working inside an IDE like Pycharm, pytest is one of the testing options. You can add it to your run configuration and then use that configuration each time you run the tests. Outside of an IDE, you can pip install pytest and then run the application using the test file as a parameter as shown here.

{% highlight python %} pytest test_my_file.py {% endhighlight %}

Using the AWS SDK

Testing with external dependencies, such as requests to external systems, always feels like a challenge to me. In a pure unit test, you remove dependencies so that you’re truly only testing your code. With external dependencies, that can be difficult because you are either required to mock the objects and pass them to your function or class or have some type of fake that replicates the expected behavior. I rarely see a fake object included in a library, which means that you’ll need to create your own. This is time consuming and can be more work than creating the code to be tested.

Working with the AWS SDK is no different. It’s an external resource that you don’t want to contact every time you run a test or at all probably. So what are your options? You can create a wrapper around the resource and mock it when needed.

{% highlight python %} class AWSWrapper: def get_resource(resource_type: str): return boto3.resource(resource_type) {% endhighlight %}

That’s certainly an option and would remove the external dependency from your unit testing. It can complicate things though because now you need to determine what every response from AWS would be, which puts a large onus on the developer to find this information. Understandably, there probably won’t be any tests as a result.

moto Library

So something else must be done then to encourage testing. One option is the motolibrary. moto intercepts calls to AWS services and allows you to make requests and receive responses as if they are from real AWS services. It can help remove barriers to testing due to its ease of use and can help increase test code coverage as a result. Below I’ll demonstrate how it can be used to test a function that uses DynamoDB.

moto is available from PyPI and can be installed with pip.

Example using DynamoDB

Code to be Tested

For this example, I’ve written a function that queries a DynamoDB table and determines if a post has been inserted into the table based on the title of the post. In this example, the title of the post is the primary key. If the entry is present, the function returns True and False if it’s not. Note that although the Count attribute is checked, there will only be one entry in the table for a given title because the title is the primary key.

If you haven’t used the DynamoDB API before, don’t confuse ScannedCount with Count. ScannedCount is the number of database entries that were scanned and is not directly related to Count. If you’d like more information about this subject, see the AWS developer guide.

Below is the function that will be tested.

{% highlight python %} def db_contains_post(table, title: str) -> bool: response = table.query( ProjectionExpression="title”, KeyConditionExpression=Key(‘title’).eq(title) )

return response['Count'] == 1

{% endhighlight %}

Setup the AWS Resource

The resource setup for this test makes up the majority of this particular test function. It might look like a lot but don’t be intimidated, it’s just setup and once you understand the paradigm, it can be reused across different AWS services. I’ll give an overview of the steps and then show the complete code at the end.

First, import the mock annotation that will be used for testing (mock_dynamodb2) along with the boto3 library. There is more than one way to use the library but the annotations are the simplest so they will be used here. {% highlight python %} import boto3 from boto3.dynamodb.conditions import Key from moto import mock_dynamodb2 {% endhighlight %}

Next, create the resource. {% highlight python %} dyanmodb = boto3.resource(‘dynamodb’) {% endhighlight %} This won’t create a DynamoDB table in AWS or in the local DynamoDB, if you’re using that. This call, like any other resource-related call will be intercepted by the framework.

Once the resource is created, create the table with the required configuration. {% highlight python %} table = dynamodb.create_table( TableName='posts’, KeySchema=[ { ‘AttributeName’: ‘title’, ‘KeyType’: ‘HASH’ # primary key } ], AttributeDefinitions=[ { ‘AttributeName’: ‘title’, ‘AttributeType’: ‘S’ },

],
ProvisionedThroughput={
    'ReadCapacityUnits': 5,
    'WriteCapacityUnits': 5
}

) {% endhighlight %} This table will mimic the table that exists in AWS so it should be configured the same way. In this example, a table called posts is created that contains a key of title.

After the table is created, populate it with data. {% highlight python %} table.put_item( Item={ ‘title’: ‘The Best Post in the World’, ‘tags’: [ ‘dogs’, ‘cats’, ‘horses’, ], ‘text’: ‘This is how you write the best post in the world…’ } ) {% endhighlight %} The items put into the table will be queryable so add what will be required for testing. The example above adds a single post that consist of a title, tags, and text.

Use the Code to be Tested

Once the setup is complete, the code under test is called. {% highlight python %} post_not_added = db_contains_post(table, ‘Some other post title’) post_added = db_contains_post(table, ‘The Best Post in the World’) {% endhighlight %}

There are two tests in this example. One using a title that was not added (post_not_added) and one where it was (post_added). Very creative names I know.

Although the table created above is used directly in the tests, your actual code would likely get a reference to a table from a DynamoDB service resource and use that instead. An example of that is shown immediately below.

{% highlight python %} dynamodb = boto3.resource(‘dynamodb’) table = dyamodb.Table(‘name_of_your_table’) {% endhighlight %}

Verify the Results

Finally, the results are verified using assert statements. The expected results are that post_not_added is False and post_added is True.

{% highlight python %} assert post_not_added is False assert post_added is True {% endhighlight %}

Complete code

{% highlight python %} import boto3 from boto3.dynamodb.conditions import Key from moto import mock_dynamodb2

def db_contains_post(table, title: str) -> bool: response = table.query( ProjectionExpression="title”, KeyConditionExpression=Key(‘title’).eq(title) )

return response['Count'] == 1

@mock_dynamodb2 def test_contains_post(): dynamodb = boto3.resource(‘dynamodb’)

table = dynamodb.create_table(
    TableName='posts',
    KeySchema=[
        {
            'AttributeName': 'title',
            'KeyType': 'HASH'  # primary key
        }
    ],
    AttributeDefinitions=[
        {
            'AttributeName': 'title',
            'AttributeType': 'S'
        },

    ],
    ProvisionedThroughput={
        'ReadCapacityUnits': 5,
        'WriteCapacityUnits': 5
    }
)

table.put_item(
    Item={
        'title': 'The Best Post in the World',
        'tags': [
            'dogs',
            'cats',
            'horses',
        ],
        'text': 'This is how you write the best post in the world...'
    }
)

post_not_added = db_contains_post(table, 'Some other post title')
post_added = db_contains_post(table, 'The Best Post in the World')

assert post_not_added is False
assert post_added is True

{% endhighlight %}

Run the Test

You can either run this file in your IDE or by using pytest directly. This file happens to be named test_lambda_handler.py so to test it, run pytest test_lambda_handler.py.

Below is an example of the results if you run the tests using pytest from the command line. It’s a fairly compact and simple result but it illustrates what successful tests produce.

<img src=”{{ “/img/pytest_successful_test.png” | absolute_url }}” alt="pytest successful test">

Summary

In this post, you’ve seen an example of how to use the moto library for unit testing Python code that uses the AWS SDK. It is a library that can help improve test coverage with minimal setup. Try it on some of your code and leave a comment below about your experience, positive or negative.

If you enjoyed this post and would like to know when more like it are available, follow us on Twitter.