9.6. Encapsulate Dependencies to Facilitate Mocking and Testing¶
In the previous item (see Item 78: “Use Mocks to Test Code with Complex Dependencies”), I showed how to use the facilities of the unittest.mock built-in module—including the Mock class and patch family of functions—to write tests that have complex dependencies, such as a database. However, the resulting test code requires a lot of boilerplate, which could make it more difficult for new readers of the code to understand what the tests are trying to verify.
One way to improve these tests is to use a wrapper object to encapsulate the database’s interface instead of passing a DatabaseConnection object to functions as an argument. It’s often worth refactoring your code (see Item 89: “Consider warnings to Refactor and Migrate Usage” for one approach) to use better abstractions because it facilitates creating mocks and writing tests. Here, I redefine the various database helper functions from the previous item as methods on a class instead of as independent functions:
>>> class ZooDatabase:
>>> ...
>>>
>>> def get_animals(self, species):
>>> ...
>>>
>>> def get_food_period(self, species):
>>> ...
>>>
>>> def feed_animal(self, name, when):
>>> ...
Now, I can redefine the do_rounds function to call methods on a ZooDatabase object:
>>> from datetime import datetime
>>>
>>> def do_rounds(database, species, *, utcnow=datetime.utcnow):
>>>
>>> now = utcnow()
>>> feeding_timedelta = database.get_food_period(species)
>>> animals = database.get_animals(species)
>>> fed = 0
>>>
>>> for name, last_mealtime in animals:
>>> if (now - last_mealtime) >= feeding_timedelta:
>>> database.feed_animal(name, now)
>>> fed += 1
>>>
>>> return fed
Writing a test for do_rounds is now a lot easier because I no longer need to use unittest.mock.patch to inject the mock into the code being tested. Instead, I can create a Mock instance to represent a ZooDatabase and pass that in as the database parameter. The Mock class returns a mock object for any attribute name that is accessed. Those attributes can be called like methods, which I can then use to set expectations and verify calls. This makes it easy to mock out all of the methods of a class:
>>> from unittest.mock import Mock
>>>
>>> database = Mock(spec=ZooDatabase)
>>> print(database.feed_animal)
>>> database.feed_animal()
>>> database.feed_animal.assert_any_call()
<Mock name='mock.feed_animal' id='140605153795472'>
I can rewrite the Mock setup code by using the ZooDatabase encapsulation:
>>> from datetime import timedelta
>>> from unittest.mock import call
>>>
>>> now_func = Mock(spec=datetime.utcnow)
>>> now_func.return_value = datetime(2019, 6, 5, 15, 45)
>>>
>>> database = Mock(spec=ZooDatabase)
>>> database.get_food_period.return_value = timedelta(hours=3)
>>> database.get_animals.return_value = [
>>> ('Spot', datetime(2019, 6, 5, 11, 15)),
>>> ('Fluffy', datetime(2019, 6, 5, 12, 30)),
>>> ('Jojo', datetime(2019, 6, 5, 12, 55))
>>> ]
Then I can run the function being tested and verify that all dependent methods were called as expected:
>>> result = do_rounds(database, 'Meerkat', utcnow=now_func)
>>> assert result == 2
>>>
>>> database.get_food_period.assert_called_once_with('Meerkat')
>>> database.get_animals.assert_called_once_with('Meerkat')
>>> database.feed_animal.assert_has_calls(
>>> [
>>> call('Spot', now_func.return_value),
>>> call('Fluffy', now_func.return_value),
>>> ],
>>> any_order=True)
Using the spec parameter to Mock is especially useful when mocking classes because it ensures that the code under test doesn’t call a misspelled method name by accident. This allows you to avoid a common pitfall where the same bug is present in both the code and the unit test, masking a real error that will later reveal itself in production:
database.bad_method_name()
>>>
Traceback ...
AttributeError: Mock object has no attribute 'bad_method_name'
If I want to test this program end-to-end with a mid-level integration test (see Item 77: “Isolate Tests from Each Other with setUp, tearDown, setUpModule, and tearDownModule”), I still need a way to inject a mock ZooDatabase into the program. I can do this by creating a helper function that acts as a seam for dependency injection. Here, I define such a helper function that caches a ZooDatabase in module scope (see Item 86: “Consider Module-Scoped Code to Configure Deployment Environments”) by using a global statement:
>>> DATABASE = None
>>>
>>> def get_database():
>>> global DATABASE
>>> if DATABASE is None:
>>> DATABASE = ZooDatabase()
>>> return DATABASE
>>>
>>> def main(argv):
>>> database = get_database()
>>> species = argv[1]
>>> count = do_rounds(database, species)
>>> print(f'Fed {count} {species}(s)')
>>> return 0
Now, I can inject the mock ZooDatabase using patch, run the test, and verify the program’s output. I’m not using a mock datetime.utcnow here; instead, I’m relying on the database records returned by the mock to be relative to the current time in order to produce similar behavior to the unit test. This approach is more flaky than mocking everything, but it also tests more surface area:
>>> import contextlib
>>> import io
>>> from unittest.mock import patch
>>> with patch('__main__.DATABASE', spec=ZooDatabase):
>>> now = datetime.utcnow()
>>>
>>> DATABASE.get_food_period.return_value = timedelta(hours=3)
>>> DATABASE.get_animals.return_value = [
>>> ('Spot', now - timedelta(minutes=4.5)),
>>> ('Fluffy', now - timedelta(hours=3.25)),
>>> ('Jojo', now - timedelta(hours=3)),
>>> ]
>>>
>>> fake_stdout = io.StringIO()
>>> with contextlib.redirect_stdout(fake_stdout):
>>> main(['program name', 'Meerkat'])
>>>
>>> found = fake_stdout.getvalue()
>>> expected = 'Fed 2 Meerkat(s)\n'
>>>
>>> assert found == expected
The results match my expectations. Creating this integration test was straightforward because I designed the implementation to make it easier to test.
9.6.1. Things to Remember¶
✦ When unit tests require a lot of repeated boilerplate to set up mocks, one solution may be to encapsulate the functionality of dependencies into classes that are more easily mocked.
✦ The Mock class of the unittest.mock built-in module simulates classes by returning a new mock, which can act as a mock method, for each attribute that is accessed.
✦ For end-to-end tests, it’s valuable to refactor your code to have more helper functions that can act as explicit seams to use for injecting mock dependencies in tests.