9.5. Use Mocks to Test Code with Complex Dependencies

Another common need when writing tests (see Item 76: “Verify Related Behaviors in TestCase Subclasses”) is to use mocked functions and classes to simulate behaviors when it’s too difficult or slow to use the real thing. For example, say that I need a program to maintain the feeding schedule for animals at the zoo. Here, I define a function to query a database for all of the animals of a certain species and return when they most recently ate:

>>> class DatabaseConnection:
>>>     ...
>>>
>>> def get_animals(database, species):
>>>     # Query the database
>>>     ...
>>>     # Return a list of (name, last_mealtime) tuples

How do I get a DatabaseConnection instance to use for testing this function? Here, I try to create one and pass it into the function being tested:

database = DatabaseConnection('localhost', '4444')

get_animals(database, 'Meerkat')

>>>
Traceback ...
DatabaseConnectionError: Not connected

There’s no database running, so of course this fails. One solution is to actually stand up a database server and connect to it in the test. However, it’s a lot of work to fully automate starting up a database, configuring its schema, populating it with data, and so on in order to just run a simple unit test. Further, it will probably take a lot of wallclock time to set up a database server, which would slow down these unit tests and make them harder to maintain.

A better approach is to mock out the database. A mock lets you provide expected responses for dependent functions, given a set of expected calls. It’s important not to confuse mocks with fakes. A fake would provide most of the behavior of the DatabaseConnection class but with a simpler implementation, such as a basic in-memory, single-threaded database with no persistence.

Python has the unittest.mock built-in module for creating mocks and using them in tests. Here, I define a Mock instance that simulates the get_animals function without actually connecting to the database:

>>> from datetime import datetime
>>> from unittest.mock import Mock
>>>
>>> mock = Mock(spec=get_animals)
>>> expected = [
>>>     ('Spot', datetime(2019, 6, 5, 11, 15)),
>>>     ('Fluffy', datetime(2019, 6, 5, 12, 30)),
>>>     ('Jojo', datetime(2019, 6, 5, 12, 45)),
>>> ]
>>> mock.return_value = expected

The Mock class creates a mock function. The return_value attribute of the mock is the value to return when it is called. The spec argument indicates that the mock should act like the given object, which is a function in this case, and error if it’s used in the wrong way.

For example, here I try to treat the mock function as if it were a mock object with attributes:

mock.does_not_exist

>>>
Traceback ...
AttributeError: Mock object has no attribute 'does_not_exist'

Once it’s created, I can call the mock, get its return value, and verify that what it returns matches expectations. I use a unique object value as the database argument because it won’t actually be used by the mock to do anything; all I care about is that the database parameter was correctly plumbed through to any dependent functions that needed a DatabaseConnection instance in order to work (see Item 55: “Use Queue to Coordinate Work Between Threads” for another example of using sentinel object instances):

>>> database = object()
>>> result = mock(database, 'Meerkat')
>>> assert result == expected

This verifies that the mock responded correctly, but how do I know if the code that called the mock provided the correct arguments? For this, the Mock class provides the assert_called_once_with method, which verifies that a single call with exactly the given parameters was made:

>>> mock.assert_called_once_with(database, 'Meerkat')

If I supply the wrong parameters, an exception is raised, and any TestCase that the assertion is used in fails:

mock.assert_called_once_with(database, 'Giraffe')

>>>
Traceback ...
AssertionError: expected call not found.
Expected: mock(<object object at 0x109038790>, 'Giraffe')
Actual: mock(<object object at 0x109038790>, 'Meerkat')

If I actually don’t care about some of the individual parameters, such as exactly which database object was used, then I can indicate that any value is okay for an argument by using the unittest.mock.ANY constant. I can also use the assert_called_with method of Mock to verify that the most recent call to the mock—and there may have been multiple calls in this case—matches my expectations:

>>> from unittest.mock import ANY
>>> mock = Mock(spec=get_animals)
>>> mock('database 1', 'Rabbit')
>>> mock('database 2', 'Bison')
>>> mock('database 3', 'Meerkat')
>>>
>>> mock.assert_called_with(ANY, 'Meerkat')

ANY is useful in tests when a parameter is not core to the behavior that’s being tested. It’s often worth erring on the side of under-specifying tests by using ANY more liberally instead of over-specifying tests and having to plumb through various test parameter expectations.

The Mock class also makes it easy to mock exceptions being raised:

class MyError(Exception):

pass

mock = Mock(spec=get_animals) mock.side_effect = MyError('Whoops! Big problem') result = mock(database, 'Meerkat')

>>>
Traceback ...
MyError: Whoops! Big problem

There are many more features available, so be sure to see help(unittest.mock.Mock) for the full range of options.

Now that I’ve shown the mechanics of how a Mock works, I can apply it to an actual testing situation to show how to use it effectively in writing unit tests. Here, I define a function to do the rounds of feeding animals at the zoo, given a set of database-interacting functions:

>>> def get_food_period(database, species):
>>>     # Query the database
>>>     ...
>>>     # Return a time delta
>>>
>>> def feed_animal(database, name, when):
>>>     # Write to the database
>>>     ...
>>>
>>> def do_rounds(database, species):
>>>     now = datetime.datetime.utcnow()
>>>     feeding_timedelta = get_food_period(database, species)
>>>     animals = get_animals(database, species)
>>>     fed = 0
>>>     for name, last_mealtime in animals:
>>>         if (now - last_mealtime) > feeding_timedelta:
>>>             feed_animal(database, name, now)
>>>             fed += 1
>>>
>>>     return fed

The goal of my test is to verify that when do_rounds is run, the right animals got fed, the latest feeding time was recorded to the database, and the total number of animals fed returned by the function matches the correct total. In order to do all this, I need to mock out datetime.utcnow so my tests have a stable time that isn’t affected by daylight saving time and other ephemeral changes. I need to mock out get_food_period and get_animals to return values that would have come from the database. And I need to mock out feed_animal to accept data that would have been written back to the database.

The question is: Even if I know how to create these mock functions and set expectations, how do I get the do_rounds function that’s being tested to use the mock dependent functions instead of the real versions? One approach is to inject everything as keyword-only arguments (see Item 25: “Enforce Clarity with Keyword-Only and Positional-Only Arguments”):

>>> def do_rounds(database, species, *,
>>>               now_func=datetime.utcnow,
>>>               food_func=get_food_period,
>>>               animals_func=get_animals,
>>>               feed_func=feed_animal):
>>>     now = now_func()
>>>     feeding_timedelta = food_func(database, species)
>>>     animals = animals_func(database, species)
>>>     fed = 0
>>>
>>>     for name, last_mealtime in animals:
>>>         if (now - last_mealtime) > feeding_timedelta:
>>>             feed_func(database, name, now)
>>>             fed += 1
>>>
>>>     return fed

To test this function, I need to create all of the Mock instances upfront and set their expectations:

>>> from datetime import timedelta
>>>
>>> now_func = Mock(spec=datetime.utcnow)
>>> now_func.return_value = datetime(2019, 6, 5, 15, 45)
>>>
>>> food_func = Mock(spec=get_food_period)
>>> food_func.return_value = timedelta(hours=3)
>>>
>>> animals_func = Mock(spec=get_animals)
>>> animals_func.return_value = [
>>>     ('Spot', datetime(2019, 6, 5, 11, 15)),
>>>     ('Fluffy', datetime(2019, 6, 5, 12, 30)),
>>>     ('Jojo', datetime(2019, 6, 5, 12, 45)),
>>> ]
>>>
>>> feed_func = Mock(spec=feed_animal)

Then, I can run the test by passing the mocks into the do_rounds function to override the defaults:

>>> result = do_rounds(
>>>     database,
>>>     'Meerkat',
>>>     now_func=now_func,
>>>     food_func=food_func,
>>>     animals_func=animals_func,
>>>     feed_func=feed_func)
>>>
>>> assert result == 2

Finally, I can verify that all of the calls to dependent functions matched my expectations:

>>> from unittest.mock import call
>>>
>>> food_func.assert_called_once_with(database, 'Meerkat')
>>>
>>> animals_func.assert_called_once_with(database, 'Meerkat')
>>>
>>> feed_func.assert_has_calls(
>>>     [
>>>         call(database, 'Spot', now_func.return_value),
>>>         call(database, 'Fluffy', now_func.return_value),
>>>     ],
>>>     any_order=True)

I don’t verify the parameters to the datetime.utcnow mock or how many times it was called because that’s indirectly verified by the return value of the function. For get_food_period and get_animals, I verify a single call with the specified parameters by using assert_called_once_with. For the feed_animal function, I verify that two calls were made—and their order didn’t matter—to write to the database using the unittest.mock.call helper and the assert_has_calls method.

This approach of using keyword-only arguments for injecting mocks works, but it’s quite verbose and requires changing every function you want to test. The unittest.mock.patch family of functions makes injecting mocks easier. It temporarily reassigns an attribute of a module or class, such as the database-accessing functions that I defined above. For example, here I override get_animals to be a mock using patch:

>>> from unittest.mock import patch
>>>
>>> print('Outside patch:', get_animals)
>>>
>>> with patch('__main__.get_animals'):
>>>     print('Inside patch: ', get_animals)
>>>
>>> print('Outside again:', get_animals)
Outside patch: <function get_animals at 0x7f5e21790310>
Inside patch:  <MagicMock name='get_animals' id='140042256426656'>
Outside again: <function get_animals at 0x7f5e21790310>

patch works for many modules, classes, and attributes. It can be used in with statements (see Item 66: “Consider contextlib and with Statements for Reusable try/finally Behavior”), as a function decorator (see Item 26: “Define Function Decorators with functools.wraps”), or in the setUp and tearDown methods of TestCase classes (see Item 76: “Verify Related Behaviors in TestCase Subclasses”). For the full range of options, see help(unittest.mock.patch).

However, patch doesn’t work in all cases. For example, to test do_rounds I need to mock out the current time returned by the datetime.utcnow class method. Python won’t let me do that because the datetime class is defined in a C-extension module, which can’t be modified in this way:

fake_now = datetime(2019, 6, 5, 15, 45)

with patch('datetime.datetime.utcnow'):

datetime.utcnow.return_value = fake_now

>>>
Traceback ...
TypeError: can't set attributes of built-in/extension type
➥'datetime.datetime'

To work around this, I can create another helper function to fetch time that can be patched:

>>> def get_do_rounds_time():
>>>     return datetime.datetime.utcnow()
>>>
>>> def do_rounds(database, species):
>>>     now = get_do_rounds_time()
>>>     ...
>>>
>>> with patch('__main__.get_do_rounds_time'):
>>>     ...

Alternatively, I can use a keyword-only argument for the datetime.utcnow mock and use patch for all of the other mocks:

>>> def do_rounds(database, species, *, utcnow=datetime.utcnow):
>>>     now = utcnow()
>>>     feeding_timedelta = get_food_period(database, species)
>>>     animals = get_animals(database, species)
>>>     fed = 0
>>>
>>>     for name, last_mealtime in animals:
>>>         if (now - last_mealtime) > feeding_timedelta:
>>>             feed_func(database, name, now)
>>>             fed += 1
>>>
>>>     return fed

I’m going to go with the latter approach. Now, I can use the patch.multiple function to create many mocks and set their expectations:

>>> from unittest.mock import DEFAULT
>>>
>>> with patch.multiple('__main__',
>>>                     autospec=True,
>>>                     get_food_period=DEFAULT,
>>>                     get_animals=DEFAULT,
>>>                     feed_animal=DEFAULT):
>>>     now_func = Mock(spec=datetime.utcnow)
>>>     now_func.return_value = datetime(2019, 6, 5, 15, 45)
>>>     get_food_period.return_value = timedelta(hours=3)
>>>     get_animals.return_value = [
>>>         ('Spot', datetime(2019, 6, 5, 11, 15)),
>>>         ('Fluffy', datetime(2019, 6, 5, 12, 30)),
>>>         ('Jojo', datetime(2019, 6, 5, 12, 45))
>>>     ]

With the setup ready, I can run the test and verify that the calls were correct inside the with statement that used patch.multiple:

result = do_rounds(database, 'Meerkat', utcnow=now_func) assert result == 2

food_func.assert_called_once_with(database, 'Meerkat') animals_func.assert_called_once_with(database, 'Meerkat') feed_func.assert_has_calls( [

call(database, 'Spot', now_func.return_value), call(database, 'Fluffy', now_func.return_value),

], any_order=True)

The keyword arguments to patch.multiple correspond to names in the __main__ module that I want to override during the test. The DEFAULT value indicates that I want a standard Mock instance to be created for each name. All of the generated mocks will adhere to the specification of the objects they are meant to simulate, thanks to the autospec=True parameter.

These mocks work as expected, but it’s important to realize that it’s possible to further improve the readability of these tests and reduce boilerplate by refactoring your code to be more testable (see Item 79: “Encapsulate Dependencies to Facilitate Mocking and Testing”).

9.5.1. Things to Remember

✦ The unittest.mock module provides a way to simulate the behavior of interfaces using the Mock class. Mocks are useful in tests when it’s difficult to set up the dependencies that are required by the code that’s being tested.

✦ When using mocks, it’s important to verify both the behavior of the code being tested and how dependent functions were called by that code, using the Mock.assert_called_once_with family of methods.

✦ Keyword-only arguments and the unittest.mock.patch family of functions can be used to inject mocks into the code being tested.