515hikaru
Posted on May 26, 2022
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='user@example.com')
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.
- This function returns correctly user attributes.
- This function returns
None
whenadmin_get_user
raises ClientError withUserNotFoundException
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': 'user@example.com'}
]
})
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': 'user@example.com'}
]
)
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!
Posted on May 26, 2022
Join Our Newsletter. No Spam, Only the good stuff.
Sign up to receive the latest update from our blog.