Better Python Testing: How to Mock AWS [video]

, HipChat DevOps Engineer | March 08, 2017

At Atlassian, my fellow engineers and I often find ourselves creating tools that rapidly become crucial to our own workflows. Because of this, the practice of writing tests to ensure consistent behavior is almost second nature for us.

We’re aware that this is not the case for everyone.

We also know first hand that tests are often the first corner to be cut. To help avoid unnecessary heartburn, headache, and chaos, we’ve put together this video and post to help you cut your teeth on Python unit testing, and on mocking out external services. Think of this as a gentle, handheld introduction to keeping the promises your code makes.

 

Writing tests isn’t always fun, but it’s an essential part of making tools that people can depend on. We’ve gotten several papercuts while writing tests that require interaction with AWS Resources, and we thought it might be helpful to share a minimal Python package skeleton which demonstrates how to set up and write tests in Python, and how to mock AWS endpoints in your tests. The benefits of this are pretty substantial:

  • Your tests will run faster, as the local mocked endpoint will respond faster than an endpoint in some remote data center.
  • It reduces the chances of erroneous failure from network instability.
  • It increases the safety of your testing, as it will be much harder to ever impact a live environment with them.
  • It increases the security of your environment, as you don’t have to distribute keys along with your tests somehow.
  • It reduces the expense of testing, as you don’t have to interact with or maintain resources that cost money and brainpower to have exist.

The code samples and test cases are available in this Bitbucket repository.

3rd party Python packages involved

The external packages used in this example are pretty minimal:

  • Unittest is the library that provides Python’s standard testing framework and will be used in each of our test cases.
  • Coverage is a 3rd party package that helps us visualize and measure our project’s test case coverage.
  • Boto is Amazon’s Python SDK.
  • Moto is a fantastic tool for mocking boto calls.

Setup

Create a Python Package

To make the scaffolding a little easier to build and use, we put the project into a Python package. If you’re not familiar with building a package, the process can be a little confusing, but it’s just a bit of Python and some layout rules.

We’re not gonna dive into the nitty-gritty of packaging Python code, but we’ll outline the needful here. If you want a good Python packaging tutorial, we’d offer the distutils setupscript or The Hitchhiker’s Guide to Packaging as a starting point.

First, we add a setup.py file, which is responsible for telling Python important data about our package.

Inside that, we import some functions from setuptools and make a call to the setup function, which is responsible for giving our package a name, version, and a few other bits.

  • The packages entry is a list of directories which contain our Python modules.
  • The test_suite entry is the directory containing our test cases.
  • install_requires is a list of each package needed to run our project.
  • Unsurprisingly, tests_require is a list of each package needed to test our project.
  • Lastly, entry_points allows us to create system wide callables to our Python functions.

Module Design

The more concise your function or module, the easier it will be to document and test. This will also be helpful when trying to read and comprehend it later. In this code, we’ve made the functions exceptionally small to make it easier to follow along.

We use boto3 to interface with two popular faces of the Amazon Web Services offerings:

  • Simple Storage Service (s3)
  • Domain Name System (route53)

Each of our interfaces has its own Python module in the project’s package directory. In this case, that’s the sample_project directory:

Each module includes a main function that’s declared as an entry_point within setup.py:

The s3 command will print out each bucket and its content:

$ s3
[ static ]
 => style.css
 => style.js

The route command prints out each domain and its records:

$ route53
[ example.com. ]
 => www.example.com.
 => blog.example.com.

Test Setup

Now that we have some modules that do stuff, we can write corresponding tests to validate that our code does what we expect.

There’s a little bit of required scaffolding to make this easy to do. First, we declare that our test code lives in the ‘tests’ directory. We do this in setup.py, specifically by setting the test_suite key to point at the the tests directory:

test_suite='tests'

For each module in our project, there’s a corresponding test module. For grokability, we’ve named these test modules with the suffix _test:

These modules contain a test for each function found in the code they test. Each test function’s name is prefixed with test_ for clarity.

Writing Test Cases

When writing unittests, and mocks for those tests, you’re essentially describing the scenario where our module’s behavior will be tested. They usually follow the general flow of:

1. ‘Under these circumstances’

This is handled mostly by the setUp method. This is where we set the stage for our test case.

2. ‘the thing being tested’

This is what this specific test relates to. The practice of making short tests against small functions helps to keep your code easy to understand.

3. ‘should behave in this fashion’

You’ll mostly be using unittest’s assertion methods here. You’ll call the module, and then verify that the results returned match what you expect.

The sample_project/s3.py‘s module’s tests are stored in s3_tests.py. Open both sample_project/s3.py and tests/s3_test.py in your favorite editor/ide, and lets quickly walk through it.

s3.py defines a function named list_s3_buckets. The coresponding companion test, s3_test.py defines a method named test_list_s3_buckets.

First, we create a new class sub-classed from unittest.TestCase, and create a setUp method. This method will execute before each test case, and is the proper place to store some constant variables. In this case, we specify and assign the variables for our mocked bucket, object key, and value which we will use in the subsequent test cases.

Next, we create a new method named _moto_setup. Notice that this method is declared with mock_s3. This silently redirects each call utilizing boto to our mock, so rather than ever talking to AWS, we hit the mocked references moto provides.

In _moto_setup we use the get_client function to get an s3 client. We then simulate bucket creation and file upload using the variables stored in setUp.

Without performing the setUp, our mocked s3 environment would contain no buckets or objects. Sometimes, this might be what you want, but not in this case.

Test Cases

OK! We are now ready to write some test cases. We take each of the project’s functions, and create a test case for them, prefixed with test_. This is the naming convention/pattern unittest expects for test cases.

In test_get_client we using a variable to hold the return of get_client. Then we use unittest‘s method assertEquals to verify the client’s endpoint.host is that of the expected s3 url. if it is, this test will pass.

Again, notice that test_list_s3_buckets is decorated with mocks3, this will force boto calls to hit our mocked environment.

Within the function, before utilizing any project components, we invoke our moto_setup method to populate our s3 mock with a bucket and object.

Now that the s3 env is setup, we’ll call list_s3_buckets function and store the output a list named buckets.

Using unittest‘s assertTrue method we can verify self.bucket, which is a constant variable defined in setUp is in returned list of buckets. If it is, this test case will pass.

test_list_s3_objects will again decorate with mocks3, then run moto_setup to setup our mock s3 env.

We can now call the list_s3_objects function using the self.bucket constant we defined in setUp, and store the results in a list named objects.

To validate that it behaves as anticipated, we use unittest‘s assertTrue method to verify that the self.key constant defined in setUp is within the returned list of objects. If it is, this test case will pass.

test_read_s3_object will again decorate with mocks3, then run moto_setup to setup our mock s3 env. We can now call the read_s3_object function using the self.bucket and self.key constants we defined in setUp, and store the results. Next we use unittest‘s assertTrue method to verify the content is the expected self.value constant variable. If it is, this test case will pass.

Last, but not least, we have test_main. First…yep, you guessed it. We decorate with mocks3, then…you’re right! We run moto_setup to setup the s3 mock.

Since this function prints output to stdout rather than returning a value, it’s a little trickier to test. We redefine stdout as a StringIO object. Then we call the project’s main function, storing the stdout data as a value in the StringIO object. This permits us to use unittest‘s assertTrue method to check that the self.bucket and self.key constants are in the stored results, and evaluate the stylized formatting. If they are as expected, this test case will pass.

Coverage

Plain text output

The usual way of running your newly crafted tests is via

python setup.py test

This is a great way to quickly see if our tests pass. What it does not do is give us information on how well the tests cover our project. This is where the coverage package comes in. To get information about how well your tests exercise your codebase, we can add coverage as a test requirement in our project so that it’s installed when we install our package.

The interface to coverage is a shell executable with the elusively obvious name coverage. To use it, simply run the

coverage

command from within your venv. The documentation for coverage can be found here.

With coverage, we are able to generate a report on the parts of our codebase that is or isn’t covered with the tests we’ve written. The report gives us useful information like the path to the module, percentage of module covered, and any lines not covered with the existing tests.

HTML coverage output

To make code coverage each nicer, coverage can output the report in html. This looks very similar to the command line report, but is more interactive, and easier to integrate into web-based CI tools.

Additionally, if any of our modules have missing bits, coverage will highlight the untested code in red blocks letting you know these lines were not tested within the tests that were run.

Gaps in coverage usually result in conversations about adding additional test cases or adapting existing cases to cover as much of my code as possible.

We hope this helps you write test cases for Python, and better understand how to mock your interactions with 3rd party services.

Have questions or feedback? Reach out on Twitter. We’d love to hear your input!