This guide provides an introduction to writing Python unit tests for your AWS cloud infrastructure. You should have at least a basic working knowledge of Python, AWS’ Python SDK (Boto3), and unit testing in general.
I’ll admit, testing is one of the things that I enjoy the least about being a developer. However, it’s so important to know how to do since it is typically a requirement in most enterprises. It also helps you catch potential defects from getting introduced into your codebase reducing time spent on fixing it later on.
Prerequisites
- Python3
- Pytest
- Boto3
- Botocore
- Text editor (i.e. VS Code, Sublime, etc.)
- Access to an AWS account
Boto3 vs Botocore
Boto3 is the official AWS SDK for Python which configures and manages services in AWS, programmatically. It is the primary module used to interact with AWS services.
Botocore is the foundation Boto3 was built upon. It is most often used in unit testing since there are modules that allows us to easily stub requests. We’ll see more about that here shortly.
Test Case
For our example, we will write a unit test for a Python class that returns a response from the Boto3 get_caller_identity method. Our unit test will be successful if the returned response matches the JSON object below:
{
"UserId": "testuser",
"Account": "012345678901",
"Arn": "arn:aws:iam::012345678901:user/testuser"
}
This is a very simple method to work with and is the reason why I chose it for this example. At the end of this guide you will know how to write a basic unit test for similar test cases utilizing stubbing, patching, and mocking.
Let’s start out by creating a Python class called, “my_module.py” and add a line to import our boto3 module. Next, we’ll create two functions. The first will return a sts client. Moving this to its own function allows us to write our unit tests more easily.
The second will serve as our function that will be directly called by our unit tests and will return the response received from the get_caller_identity call. This will provide us with something to test for. Executing this script should return a response similar to the example above.
Once finished your script should look something like this:
###
"""
Contents of my_module.py
"""
import boto3
def get_sts_client():
"""
Returns an sts boto3 client
"""
return boto3.client('sts')
def main():
"""
Main function for our module
"""
sts_client = get_sts_client()
sts_response = sts_client.get_caller_identity()
return sts_response
###
If you don’t have a default region specified in your environment you may receive a ClientError when attempting to run your script. Add the region to your sts client if you experience that error.
###
boto3.client('sts', 'us-east-1')
###
Writing our Unit Tests
Now that we’re done with our my_module.py script the next step is to write our unit tests in another file called, “test_my_module.py”. For simplicity sake, create the file in the same directory as the original file.
###
"""
Contents of test_my_module.py
"""
#Create your import statements here
from unittest.mock import patch
from botocore.stub import Stubber
import botocore
from botocore.session import Session
import my_module #Ensure this matches the file name created earlier
"""
Creating global variable here will allow you to reuse
this in multiple tests. Used as the return value when we
patch our get_sts_client function.
"""
session = Session()
STS_CLIENT = session.create_client('sts')
"""
This will serve as both our sts
response test fixture and expected result.
As a best practice, this should be created
as a separate file in another directory
where other test fixtures are located.
"""
sts_get_caller_identity_response = {
"UserId": "testid",
"Account": "012345678901",
"Arn": "arn:aws:iam::012345678901:user/testid"
}
"""
Note: Each of the three functions below does the same
thing but just using different methods.
"""
@patch.object(my_module, 'get_sts_client', return_value=STS_CLIENT)
def test_main_01(get_sts_client):
"""
Method 1 - Patch decorators without using stubber context manager
"""
STS_STUBBER = Stubber(STS_CLIENT)
STS_STUBBER.add_response('get_caller_identity', sts_get_caller_identity_response)
STS_STUBBER.activate()
my_module_result = my_module.main()
STS_STUBBER.assert_no_pending_responses()
assert(my_module_result == sts_get_caller_identity_response)
STS_STUBBER.deactivate()
@patch.object(my_module, 'get_sts_client', return_value=STS_CLIENT)
def test_main_02(get_sts_client):
"""
Method 2 - Patch decorators using stubber context manager
"""
with Stubber(STS_CLIENT) as STS_STUBBER:
STS_STUBBER.add_response('get_caller_identity', sts_get_caller_identity_response)
my_module_result = my_module.main()
STS_STUBBER.assert_no_pending_responses()
assert(my_module_result == sts_get_caller_identity_response)
def test_main_03():
"""
Method 3 - Patching class without stubber context manager
"""
with patch('my_module.get_sts_client') as mock:
mock_instance = mock.return_value
mock_instance.get_caller_identity.return_value = sts_get_caller_identity_response
result = my_module.main()
assert result == sts_get_caller_identity_response
###
That’s it! If you run pytest on your test_my_module.py, it should return three successful tests.
I know I mentioned having access to an AWS account as a requirement at the beginning of this article. However, if done correctly, you should be able to run your unit tests without any internet connection at all. If you receive ‘NoCredentialsError: Unable to locate credentials’, it means your unit tests are actually making live calls out to AWS and something isn’t being stubbed.
Stubbing Multiple Calls
Another thing to keep in mind is that this example only tests for a single Boto3 call. Chances are that you’ll have more than one when you are writing your tests. To accomplish that, you need to have a separate ‘add_response’ method for each call that is executed in your class and paired with the appropriate test fixture response. The order of these responses must also match what you’re testing. Otherwise, you’ll receive an error message stating the test fixture does not match the corresponding Boto3 method call.
###
sts_get_caller_identity_response = {
"UserId": "testuser",
"Account": "012345678901",
"Arn": "arn:aws:iam::012345678901:user/testuser"
}
sts_get_session_token = {
"Credentials": {
"AccessKeyId":"string",
"SecretAccessKey":"string",
"SessionToken":"string",
"Expiration":20990101
}
}
STS_STUBBER.add_response('get_caller_identity', sts_get_caller_identity_response)
STS_STUBBER.add_response('get_session_token', sts_get_session_token)
###
Once you’re able to successfully stub multiple Boto3 calls, the next step would be knowing how to handle exceptions.