Spying on instance methods with Python's mock module

Python's mock module (unittest.mock in Python 3.3 and higher) allows you to observe parameters passed to functions.

I'm a little slow, so I had to dig around to figure out how to do this.

Let's say you have a class:

class ProductionClass:

    def user_api(self, a, b=False):
        if b:
            return self._internal_api_one(a)
        else:
            return self._internal_api_two(a)

    def _internal_api_one(self, a):
        # Do something necessary with a
        pass

    def _internal_api_two(self, a):
        # Do something necessary with a
        pass

Now, let's suppose you are testing the functionality of ProductionClass, but you want to observe the parameters passed to your internal methods but still invoke those internal methods. I didn't find a lot of examples of this from my Google searches, so here is the solution using unittest.mock (or mock from PyPI if you're on Legacy Python 2.x):

import unittest.mock as mock

def test_user_api():
    inst = ProductionClass()

    with mock.patch.object(inst, '_internal_api_one',
                           wraps=inst._internal_api_one) as monkey:
         inst.user_api('foo', b=True)
         monkey.assert_called_with('foo')

    with mock.patch.object(inst, '_internal_api_two',
                           wraps=inst._internal_api_two) as monkey:
         inst.user_api('bar', b=False)
         monkey.assert_called_with('bar')

Happy hunting.