Practical Example of boto3 Stubber Classes

I'm a Python developer, using Django (with rest framework) and AWS services. Of course I use boto3 to connect with some AWS services from application.

By the way, when you add some tests for using boto3 code you have to use mock because shouldn't connect real resources when running unit tests. I had used unittest.mock at first, but I found stubber class written in official document.

Stubber Reference — botocore 1.26.7 documentation

But I didn't find practical example code, so I'm writing this post now.

Overview

First of all, The following code is how to use stubber class(I will introduce more detail next sections).

>>> import boto3
>>> from botocore.stub import Stubber
>>> client = boto3.client('cognito-idp')
>>> stubber = Stubber(client)
>>> stubber.add_response('list_users', {'Users': []})
>>> stubber.add_client_error('admin_get_user', service_error_code='UserNotFoundException')
>>> stubber.activate()
>>> client.list_users(UserPoolId='dummpy_id')
{'Users': []}
>>> client.admin_get_user(UserPoolId='dummpy_id', Username='[e-mail]')
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/hikaru/.local/lib/python3.8/site-packages/botocore/client.py", line 357, in _api_call
    return self._make_api_call(operation_name, kwargs)
  File "/home/hikaru/.local/lib/python3.8/site-packages/botocore/client.py", line 676, in _make_api_call
    raise error_class(parsed_response, operation_name)
botocore.errorfactory.UserNotFoundException: An error occurred (UserNotFoundException) when calling the AdminGetUser operation:

Initialization

You can construct client as usual. In my example, specifying cognito-idp create client that can operate AWS Cognito user, group, and so on.

>>> client = boto3.client('cognito-idp')
>>> stubber = Stubber(client)

Pass Pattern

add_response method can mock boto3 client's method. add_response receive two args, first is name of the target method, second is return value. For example, cognito-idp client has list_users method. Following code will patch client.list_users(some_args) to return {'Users': []}.

>>> stubber.add_response('list_users', {'Users': []})

Error

You can mock to raise the error. Stubber class has add_client_error method. This also receives 2 args, first is same as method name, second is error codes.

The following code is example of patching admin_get_user to raise UserNotFoundException.

>>> stubber.add_client_error('admin_get_user', service_error_code='UserNotFoundException')

Activation

Call activate() method or create with block.

stubber.activate()
client.list_users()

# or

with stubber:
    client.list_users()

If you want to know more detail, read the official reference.

Example of unit test code

As you can see, you can mock flexibly. But you might say that you don't know how to use this.

So, through the more practical example, I introduce how to use stubber class.

Sample of function

I prepare simple example to test. The function get_user is receiving email and returning attributes of the user who has same email address.

# main.py
import os

import boto3
from botocore.exceptions import ClientError

os.environ['AWS_ACCESS_KEY_ID'] = 'DUMMY_VALUE'
os.environ['AWS_SECRET_ACCESS_KEY'] = 'DUMMY_VALUE'
os.environ['AWS_DEFAULT_REGION'] = 'ap-northeast-1'

def get_user(email: str) -> dict:
    client = boto3.client('cognito-idp')

    try:
        user = client.admin_get_user(
            UserPoolId='DUMMY_USER_POOL_ID',
            Username=email,
        )
    except ClientError as error:
        if error.response['Error']['Code'] == 'UserNotFoundException':
            return None
        raise

    return user['UserAttributes']

Test Strategy

Probably, you would like to write 2 test cases at least.

1. This function returns correctly user attributes.
2. This function returns None when admin_get_user raises ClientError with UserNotFoundException code.

So, I will use to Stubber class for writing two tests.

The first case

In this example, I use unittest.mock in Python standard library.

# tests/test_main.py
import unittest
from unittest import mock
import boto3
from botocore.stub import Stubber

from main import get_user

class TestGetUser(unittest.TestCase):

    def test_get_user(self):
        client = boto3.client('cognito-idp')
        stubber = Stubber(client)
        stubber.add_response('admin_get_user', {
            'Username': 'user',
            'UserAttributes':
                [
                    {'Name': 'sub', 'Value': 'aa45403e-8ba5-42ab-ab27-78a6e9335b23'},
                    {'Name': 'email', 'Value': '[e-mail]'}
                ]
        })
        stubber.activate()
        with mock.patch('boto3.client', mock.MagicMock(return_value=client)):
            user = get_user('user')

        self.assertEqual(user,
                [
                    {'Name': 'sub', 'Value': 'aa45403e-8ba5-42ab-ab27-78a6e9335b23'},
                    {'Name': 'email', 'Value': '[e-mail]'}
                ]
        )

A notable point is with mock.patch('boto3.client', mock.MagicMock(return_value=client)):.

In this case, you have to do not only to activate stubber, but also to replace boto3.client with stubber, because boto3.client which is constructed in test method is not used in get_user function. The with block is doing such a thing.

Everything else is normal test codes.

The second case

Second case are almost same to first. In the code below, import statement and definition of TestCase class are omitted.

def test_not_found_user(self):
    client = boto3.client('cognito-idp')
    stubber = Stubber(client)
    stubber.add_client_error('admin_get_user', 'UserNotFoundException')
    stubber.activate()
    with mock.patch('boto3.client', mock.MagicMock(return_value=client)):
        user = get_user('user')

    self.assertIsNone(user)

With using add_client_error, you can make UserNotFoundError occur. So, ClientError is raised in get_user, this function returns None as you expected.

Additional Contents

The whole code of this example is here.

Conclusion

I wrote practical example of unit tests with boto3.

boto3's method returns big dictionary, so you may feel difficult to treat the response. But, in fact, all you have to do is assert only data which are related to value used in application.

Good unit tests make it easier to find bugs.

Let's use mock and write more unit tests.

Thank you for reading this article!

My original article is written in Japanese.