Derek Stegelman

Keeping it DRY with Mock and Python

This week while writing a new feature in the core of one our apps, I had to mock something that would be used in almost every single test in the test suite. I'm a huge fan of keeping things DRY so you can imagine that when I started to type things like this:

import mock

class NotificationsUtilsTest(base.NotificationTestBase, TestCase):

    def setUp(self):
        super(NotificationsUtilsTest, self).setUp()

        self.request = RequestFactory()
        self.request.user = self.user

    @mock.patch('notifications.utils.validate_eid')
    def test_get_or_create_util(self, mocked_method):
        mocked_method.return_value = True
        user = utils.get_or_create_user('randomuser')

        self.assertEqual(user.email, 'randomuser@k-state.edu')
        self.assertEqual(user.username, 'randomuser')

    @mock.patch('notifications.utils.validate_eid')
    def test_get_create_existing_user(self, mocked_method):
        mocked_method.return_value = True
        user = utils.get_or_create_user('derekst')

        self.assertEqual(user.username, 'derekst')

    @mock.patch('notifications.utils.validate_eid')
    def test_get_create_valid_eid(self, mocked_method):
        mocked_method.return_value = True
        user = utils.get_or_create_user('safe')
        self.assertIsInstance(user, User)

I wasn't too thrilled.

Mocking Via Mixin

Instead of adding patch decorators all over the place we can setup the test suite to patch inside the setup and tear down of its parent class. Instead of using mock.patch as a decorator you can setup a patch object inside of the setUp() method and set its return_value there for every method. It's important to also start and stop the patching during setUp() and tearDown().

# In tests/base.py
import mock

class NotificationTestBase(object):

    def setUp(self):
        cache.clear()

        self.validate_eid_patch = mock.patch('notifications.utils.validate_eid')
        self.validate_eid_mock = self.validate_eid_patch.start()
        self.validate_eid_mock.return_value = True

    def tearDown(self):
        self.validate_eid_patch.stop()
# In tests/test_utils.py

class NotificationsUtilsTest(base.NotificationTestBase, TestCase):

    def setUp(self):
        super(NotificationsUtilsTest, self).setUp()

        self.request = RequestFactory()
        self.request.user = self.user

    def test_get_or_create_util(self):
        user = utils.get_or_create_user('randomuser')

        self.assertEqual(user.email, 'randomuser@k-state.edu')
        self.assertEqual(user.username, 'randomuser')

    def test_get_create_existing_user(self):
        user = utils.get_or_create_user('derekst')

        self.assertEqual(user.username, 'derekst')

    def test_get_create_valid_eid(self):
        user = utils.get_or_create_user('safe')
        self.assertIsInstance(user, User)

Benefits

In the end this is a pretty simple and straightforward to add a bunch of mocking to all or most of your tests. The only big thing I don't like about this solution is that it's not very obvious for newer users to the codebase as a bunch of mocking code (which can throw tests way off) is stuck in a base class mixin and not explicit on each test.