Unit Testing in Python: Tips, Tricks, and Gotchas

When writing unit tests in Python, you'll often want to patch an entire class. In my case, I've got a function that initializes a class, and that class needs to be patched in its entirety.

Note: This example will be used throughout this guide.

I want to test a function called add_numbers:

from utils import Slack

def add_numbers(x, y):
    try:
        result = x + y
    except TypeError:
        result = 'Could not add: {0} and {1}'.format(x, y)

    slack = Slack()
    slack.send(result)

Basically it's a function that adds two numbers, and sends the response to a slack channel via an API request. That API request is handled through another class called Slack (located in utils.py):

import json
import requests

SLACK_HOOK = 'https://hooks.slack.com/services/ab/cd/ef'


class Slack:
    def __init__(self, *args, **kwargs):
        # Use the hook_url pass or default to SLACK_HOOK
        self.hook_url = kwargs.get('hook_url', SLACK_HOOK)

    def send(self, msg):
        headers = {'Content-type': 'application/json'}
        payload = {'text': msg}
        requests.post(
                self.hook_url,
                data=json.dumps(payload),
                headers=headers,
                )

When unit testing, we want to test the smallest pieces of application code in isolation. And even more importantly, we don't want to write unit tests that fire off API requests to outside services. So, we'll want to patch the Slack class when writing unit tests for the add_numbers function.

So here is how we patch the Slack class for this test.

import unittest

from add_numbers import add_numbers
from unittest import mock, TestCase


class TestAddNumbers(TestCase):
    @mock.patch('add_numbers.Slack')
    def test_calls_slack_send(self, mock_slack):
        _slack = mock_slack.return_value
        add_numbers(1, 2)

        # Assert send called
        self.assertEqual(_slack.send.call_count, 1)


if __name__ == '__main__':
    unittest.main()

From here on, I'll exclude the imports and the call to our test runner.

What I want to point out about this example is how the Slack class is patched.

First, the patch decorator is used. If you're a keen observer, you'll notice we're patching add_numbers.Slack as opposed to utils.Slack (where the class actually lives). We'll discuss this in a later step.

Notice that mock_slack is passed in as an argument to our test method. This argument is of type MagicMock, and represents the mocked class. In order to get the instantiated object, you need to called mock_slack.return_value. And with this, we have access to mocked class methods, like send.

If you're asking yourself, "Why doesn't my mock work?", you should start by making sure you're patching the correct path.

In the above example, I'm tempted to mock the Slack class by patching utils.Slack. Makes sense, right? The Slack class is located in the utils.py file, so that's where it should be patched.

class TestAddNumbers(TestCase):
    @mock.patch('utils.Slack')
    def test_calls_slack_send(self, mock_slack):
        ...

Unfortunately, not in this case.

Patch the module in the scope that it's used

The add_numbers.py module imports Slack from utils like this:

from utils import Slack

The way that Slack is imported makes the difference. The import statement will search for the Slack module and bind it to the local scope. That's why we're able to use the name Slack in our module. But it also means we need to patch Slack in the scope in which it's used (add_numbers.Slack).

class TestAddNumbers(TestCase):
    @mock.patch('add_numbers.Slack')
    def test_calls_slack_send(self, mock_slack):
        ...

If you're using class based tests inheriting from unittest.TestCase there might come a time when you want to patch a module in all of your tests. There are a few ways to do this.

For this example, we'll patch requests.post in the unit tests for our Slack class.

The most obvious way is to apply the patch at the class level.

Apply the patch at the class level

import unittest

from unittest import mock
from utils import Slack


@mock.patch('requests.post')
class TestSlack(unittest.TestCase):
    def test_send_should_call_requests_post_once(self, mock_post):
        slack = Slack()
        slack.send('hello world')
        self.assertEqual(mock_post.call_count, 1)


if __name__ == '__main__':
    unittest.main()

You'll notice that requests.post is mocked at the class level, and the mocked method is available in my unit test. The mocked method is passed in as mock_post just like it would be if we applied the patch directly to the test method. The value here is fairly obvious. We're not testing the requests.post class, nor are we testing the slack endpoint. We're simply testing our Slack class, so we should never actually send an HTTP request.

But, suppose you have a lot of modules that you need to patch for a particular test. That isn't necessarily so in our case, but there are plenty of times where we'll want to apply multiple patches to all unit test methods. The downside of applying the patch to the class is that all of our test methods will need to accept all of the arguments. This isn't a huge deal, but it can be cumbersome. Imagine if you have 30 tests and they all looked like this:

def test_some_feature(self, mock_post, mock_get, mock_get_secret, mock_get_cache, mock_db_save, mock_etc):
   ...

You could obviously clean it up:

def test_some_feature(self, *args)
    ...

But keeping tracking of the index of the mock modules you want to interact with will be gross.

Apply the patch in setUp

You can actually apply the patch in the setUp method and remove the patch in the tearDown method. These methods are called for each test method, so there is a little bit more overhead, but it might be worth it. Using this method, our test class now looks like this:

class TestSlack(unittest.TestCase):
    def setUp(self):
        self.post_patch = mock.patch('requests.post')
        self.mock_post = self.post_patch.start()

    def tearDown(self):
        self.post_patch.stop()

    def test_send_should_call_requests_post_once(self):
        slack = Slack()
        slack.send('hello world')
        self.assertEqual(self.mock_post.call_count, 1)

Often we'll find it useful to create a base class for our unit tests. This helps us avoid duplicative and messy code and provides more stability for our test suite.

Simple use case

If we simply want to run the same unit tests for a variety of classes, we can use a mixin.

class MessageTestMixin:
    def test_common_test(self):
        self.assertEqual(self.status, 'running')

class TestSlack(unittest.TestCase, MessageTestMixin):
    status = 'running'

Now, this is quite a contrived example, but you it should give a clear picture of how the test mixin works. Since the mixin doesn't inherit from TestCase, the test won't be run in the context of the mixin. It will only be run as a part of the TestSlack test suite.

Limitations

While this works in simple cases, there are various limitations to this approach. What if we wanted to include setUp and tearDown methods?

class MessageTestMixin:
    def setUp(self):
        self.status = 'running'

    def tearDown(self):
        self.status = 'not running'

    def test_common_test(self):
        self.assertEqual(self.status, 'running')


class TestSlack(unittest.TestCase, MessageTestMixin):
    pass

Uh oh. The setUp and tearDown methods were not called because the mixin doesn't inherit from TestCase. But if we make the mixin inherit from TestCase, it will run the tests for the mixin, which we don't want.

======================================================================
ERROR: test_common_test (__main__.TestSlack)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "test_mixins.py", line 15, in test_common_test
    self.assertEqual(self.status, 'running')
AttributeError: 'TestSlack' object has no attribute 'status'

----------------------------------------------------------------------
Ran 1 test in 0.001s

FAILED (errors=1)

Using a base class

Rather than a mixin, we can use a base class. It will basically be the same construction as the mixin with some small adjustments.

class MessageTestBase(unittest.TestCase):
    __test__ = False

    def setUp(self):
        self.status = 'running'

    def tearDown(self):
        self.status = 'not running'

    def test_common_test(self):
        self.assertEqual(self.status, 'running')


class TestSlack(MessageTestMixin):
    __test__ = True

Notice the differences. Our mixin becomes a base class that inherits from TestCase. we include __test__ = False in the base class to prevent the test runner from executing tests in this class. Then the child class only inherits from MessageTestBase and includes __test__ = True to instruct the test runner to run our tests.

..
----------------------------------------------------------------------
Ran 2 tests in 0.001s

OK

And there we go!

In some cases, the function we're testing will make multiple calls to another function. There are various ways we'll want to test this.

To continue with the Slack example, imagine it's really important that our API call is successful. So we wrap the request in a loop to retry on failure.

class Slack:
    def __init__(self, *args, **kwargs):
        # Use the hook_url pass or default to SLACK_HOOK
        self.hook_url = kwargs.get('hook_url', SLACK_HOOK)

    def send(self, msg):
        headers = {'Content-type': 'application/json'}
        payload = {'text': msg}
        while True:
            response = requests.post(
                    self.hook_url,
                    data=json.dumps(payload),
                    headers=headers,
                    )
            if response.status_code != 200:
                time.sleep(1)

Note: To do this in real life, you'll need to add alarms or other such mechanisms to ensure your script doesn't run forever.

In this case, requests.post is called multiple times. So let's simulate a scenario where the request fails twice before it succeeds. To do this we'll need a MockResponse class.

class MockResponse:
    def __init__(self, status_code, text=''):
        self.status_code = status_code
        self.text = text

And our test will look like this:

    @mock.patch('requests.post')
    def test_send_should_retry_until_successful(self, mock_post):
        mock_post.side_effect = [
                MockResponse(status_code=503),
                MockResponse(status_code=503),
                MockResponse(status_code=200),
                ]

        slack = Slack()
        slack.send('hello world')

        self.assertEqual(mock_post.call_count, 3)

We can pass a list to the side_effect attribute on our mocked object, and this will cause the mocked method to return each item in the list as it is called multiple times.

If you've come across any unique testing scenarios that I should include in this guide, please comment below. I'd be happy to learn from you.

Also, check out a few of the resources I find myself visiting routinely regarding this subject.

• Python Mocking Cookbook
• unittest.mock - Getting Started