Sometimes, we want to test code that requires an object be supplied that is either expensive or difficult to construct. While this may mean your API needs rethinking to have a more testable interface (which typically means a more usable interface), we sometimes find ourselves writing test code that has a ton of boilerplate to set up objects that are only incidentally related to the code under test.

For example, imagine we have some code that keeps track of flight statuses in a key-value store (such as redis or memcache) such that we can store the timestamp and the most recent status. A basic version of such code might look like this:

import datetime
import redis

class FlightStatusTracker:
    ALLOWED_STATUSES = {'CANCELLED', 'DELAYED', 'ON TIME'}

    def __init__(self):
        self.redis = redis.StrictRedis()

    def change_status(self, flight, status):
        status = status.upper()
        if status not in self.ALLOWED_STATUSES:
            raise ValueError(
                   "{} is not a valid status".format(status))

        key = "flightno:{}".format(flight)
        value = "{}|{}".format(
            datetime.datetime.now().isoformat(), status)
        self.redis.set(key, value)

There are a lot of things we ought to test in that change_status method. We should check that it raises the appropriate error if a bad status is passed in. We need to ensure that it converts statuses to uppercase. We can see that the key and value have the correct formatting when the set() method is called on the redis object.

One thing we don't have to check in our unit tests, however, is that the redis object is properly storing the data. This is something that absolutely should be tested in integration or application testing, but at the unit test level, we can assume that the py-redis developers have tested their code and that this method does what we want it to. As a rule, unit tests should be self-contained and not rely on the existence of outside resources, such as a running Redis instance.

Instead, we only need to test that the set() method was called the appropriate number of times and with the appropriate arguments. We can use Mock() objects in our tests to replace the troublesome method with an object we can introspect. The following example illustrates the use of mock:

from unittest.mock import Mock
import py.test
def pytest_funcarg__tracker():
    return FlightStatusTracker()

def test_mock_method(tracker):
    tracker.redis.set = Mock()
    with py.test.raises(ValueError) as ex:
        tracker.change_status("AC101", "lost")
    assert ex.value.args[0] == "LOST is not a valid status"
    assert tracker.redis.set.call_count == 0

This test, written using py.test syntax, asserts that the correct exception is raised when an inappropriate argument is passed in. In addition, it creates a mock object for the set method and makes sure that it is never called. If it was, it would mean there was a bug in our exception handling code.

Simply replacing the method worked fine in this case, since the object being replaced was destroyed in the end. However, we often want to replace a function or method only for the duration of a test. For example, if we want to test the timestamp formatting in the mock method, we need to know exactly what datetime.datetime.now() is going to return. However, this value changes from run to run. We need some way to pin it to a specific value so we can test it deterministically.

Remember monkey-patching? Temporarily setting a library function to a specific value is an excellent use of it. The mock library provides a patch context manager that allows us to replace attributes on existing libraries with mock objects. When the context manager exits, the original attribute is automatically restored so as not to impact other test cases. Here's an example:

from unittest.mock import patch
def test_patch(tracker):
    tracker.redis.set = Mock()
    fake_now = datetime.datetime(2015, 4, 1)
    with patch('datetime.datetime') as dt:
        dt.now.return_value = fake_now
        tracker.change_status("AC102", "on time")
    dt.now.assert_called_once_with()
    tracker.redis.set.assert_called_once_with(
        "flightno:AC102",
        "2015-04-01T00:00:00|ON TIME")

In this example, we first construct a value called fake_now, which we will set as the return value of the datetime.datetime.now function. We have to construct this object before we patch datetime.datetime because otherwise we'd be calling the patched now function before we constructed it!

The with statement invites the patch to replace the datetime.datetime module with a mock object, which is returned as the value dt. The neat thing about mock objects is that any time you access an attribute or method on that object, it returns another mock object. Thus when we access dt.now, it gives us a new mock object. We set the return_value of that object to our fake_now object; that way, whenever the datetime.datetime.now function is called, it will return our object instead of a new mock object.

Then, after calling our change_status method with known values, we use the mock class's assert_called_once_with function to ensure that the now function was indeed called exactly once with no arguments. We then call it a second time to prove that the redis.set method was called with arguments that were formatted as we expected them to be.

The previous example is a good indication of how writing tests can guide our API design. The FlightStatusTracker object looks sensible at first glance; we construct a redis connection when the object is constructed, and we call into it when we need it. When we write tests for this code, however, we discover that even if we mock out that self.redis variable on a FlightStatusTracker, the redis connection still has to be constructed. This call actually fails if there is no Redis server running, and our tests also fail.

We could solve this problem by mocking out the redis.StrictRedis class to return a mock in a setUp method. A better idea, however, might be to rethink our example. Instead of constructing the redis instance inside__init__, perhaps we should allow the user to pass one in, as in the following example:

    def __init__(self, redis_instance=None):
        self.redis = redis_instance if redis_instance else redis.StrictRedis()

This allows us to pass a mock in when we are testing, so the StrictRedis method never gets constructed. However, it also allows any client code that talks to FlightStatusTracker to pass in their own redis instance. There are a variety of reasons they might want to do this. They may have already constructed one for other parts of their code. They may have created an optimized implementation of the redis API. Perhaps they have one that logs metrics to their internal monitoring systems. By writing a unit test, we've uncovered a use case that makes our API more flexible from the start, rather than waiting for clients to demand we support their exotic needs.

This has been a brief introduction to the wonders of mocking code. Mocks are part of the standard unittest library since Python 3.3, but as you see from these examples, they can also be used with py.test and other libraries. Mocks have other more advanced features that you may need to take advantage of as your code gets more complicated. For example, you can use the spec argument to invite a mock to imitate an existing class so that it raises an error if code tries to access an attribute that does not exist on the imitated class. You can also construct mock methods that return different arguments each time they are called by passing a list as the side_effect argument. The side_effect parameter is quite versatile; you can also use it to execute arbitrary functions when the mock is called or to raise an exception.

In general, we should be quite stingy with mocks. If we find ourselves mocking out multiple elements in a given unit test, we may end up testing the mock framework rather than our real code. This serves no useful purpose whatsoever; after all, mocks are well-tested already! If our code is doing a lot of this, it's probably another sign that the API we are testing is poorly designed. Mocks should exist at the boundaries between the code under test and the libraries they interface with. If this isn't happening, we may need to change the API so that the boundaries are redrawn in a different place.