Mocking in Python ⋆ integralist

Introduction

Mocking resources when writing tests in Python can be confusing if you’re unfamiliar with doing such things. In this post I am going to cover various aspects of mocking code, which will hopefully be a useful resource for those who are a bit stuck.

Note: in the code examples I’m using pytest, but for the most part that shouldn’t matter.

unittest.mock or mock

In order to ‘mock’ a resource we’ll first need the mock module, and this is our first stumbling block: which version do we need? i.e. there’s two and they both look to be official (mock and unittest.mock).

The mock module is a backwards compatible library you can download from PyPy, where as unittest.mock is the same thing but only compatible with the version of Python you’re using.

So in almost all cases you’ll want to import it like so:

import unittest.mock as mock

For more examples, see this reference guide

Decorator

The most common way to mock resources is to use a Python decorator around your test function:

@mock.patch("thing")
def test_stuff(mock_thing):
    mock_thing.return_value = 123

In this case, what we’re patching (thing) can be a variable or a function.

When you do this you’ll need to pass an argument to your function (you can name it whatever you want †) which will be a MagicMock.

This means if you don’t do anything else, then calls to thing will (in the example above at least) result in the value 123 being returned.

† convention is to name the variable mock_<noun>.

If you’re mocking multiple things then you’ll stack the mock decorators ontop of each other, and pass them along in order to the test function:

@mock.patch("third")
@mock.patch("second")
@mock.patch("first")
def test_stuff(mock_first, mock_second, mock_third):
    ...

Resource location

It’s important to know that when mocking you should specify the location of the resource to be mocked, relevant to where it’s imported. This is best explained by way of example…

Imagine I have a module app.foo and within that module I import another dependency like so:

from app.bar import thing

You might think that when you call mock.patch that you pass it a reference to the resource like app.bar.thing. That would only be relevant if the resource was being called with that full path within the app.foo module (e.g. if app.foo called app.bar.thing(...)).

If the full namespace path isn’t referenced, which it isn’t in the above example (notice we import just the thing resource). It means we need to specify the reference namespace to mock as where it’s imported:

@mock.patch('app.foo.thing')

So even though thing exists within app.bar we specify app.foo.thing as app.foo is where we’ve imported it for use. This catches people out all the time.

Mock return_value vs side_effect

If your function has a try/except around it, then you can use side_effect to cause the calling of the function to trigger an Exception as the returned value:

@mock.patch('app.aws.sdk.confirm_sign_up', side_effect=Exception('whoops'))

Note: if you had used return_value=Exception('whoops') then the mock would return the string representation of the Exception rather than raising an exception like side_effect does.

Otherwise if you just need a static value returned, so it’s evaluated at the time it’s defined (not when it’s called), then you can use return_value instead:

@mock.patch('app.security.secret_hash', return_value='###')

Mock Nested Calls

Calling a property on a mock returns another mock, so in order to mock very specific properties you’ll need to nest your return_value or side_effect:

m = mock.MagicMock()
m.return_value.get.side_effect = [1, 2]
m.return_value.post.return_value = 'foo'

x = m()

x.get()   # 1
x.post()  # foo
x.get()   # 2

Things can get a little more confusing when you want to verify a specific nested method on a mocked object was called:

import unittest.mock as mock

from tornado.ioloop import IOLoop

@mock.patch("__main__.IOLoop")
def foo(mock_ioloop):
    IOLoop.current()
    mock_ioloop.current.assert_called()  # will fail if assertion isn't True

    IOLoop.current().start()
    mock_ioloop.current().start.assert_called()  # will fail if assertion isn't True

The reason this can get more complicated is due to how a mock will return a new mock when accessing a property on a mock:

import unittest.mock as mock

from tornado.httpserver import HTTPServer

@mock.patch("__main__.HTTPServer")
def bar(mock_httpserver):
    server = HTTPServer()

    server.listen(8080)
    HTTPServer().listen(123)

    mock_httpserver.assert_called()
    mock_httpserver().listen.assert_called_with(8080)

The above code will error:

AssertionError: expected call not found.
Expected: listen(8080)
Actual: listen(123)

You’ll need to make sure you assert the mock at the right time:

import unittest.mock as mock

from tornado.httpserver import HTTPServer

@mock.patch("__main__.HTTPServer")
def bar(mock_httpserver):
    server = HTTPServer()
    mock_httpserver.assert_called()
    
    server.listen(8080)
    mock_httpserver().listen.assert_called_with(8080)
    
    HTTPServer().listen(123)
    mock_httpserver().listen.assert_called_with(123)

Verify Exceptions

If we want to verify that some piece of code throws an Exception type when we need it to we can mock specific resources to throw an exception and then use pytest.raises as a context manager around the caller of our code to be verified.

In the following example our code (in app.account.confirm(...)) catches a generic Exception and re-raises it as exceptions.CognitoException.

We can catch and make assertions against this expected behaviour by first mocking the resource we want to throw an exception and get it to throw our own fake exception using the side_effect parameter.

Next we specify the exact exception type we’re expecting to be raised using pytest.raises(T):

@mock.patch('app.aws.sdk.confirm_sign_up', side_effect=Exception('whoops'))
def test_account_confirm_failure(mock_signup):
    with pytest.raises(exceptions.CognitoException) as exc_info:
        app.account.confirm(123, 'foo')
        assert True is True  # this will never be executed!
        
    assert exc_info.typename == 'CognitoException'
    assert str(exc_info.value) == 'SIGNUP_CONFIRMATION_FAILED'

Note: don’t make the mistake of putting any assertions within the with context manager. Once the Exception is raised by the function being called within the with context manager, all code after it inside the block is skipped.

Clearing lru_cache

If a function you wish to test has the functools.lru_cache decorator applied, then you’ll need to be mindful of mocking the response of that function as it’ll be cached in one test and the cached result will be returned when calling the function again to test some other behaviour (and might likely confuse you when you see the unexpected response).

To fix this issue is very easy because lru_cache provides additional functions when decoratoring your functions, it provides:

cache_info
cache_clear

The latter (cache_clear) is what you would need to call. This is demonstrated below:

@lru_cache(5)
def foo():
    print('Executing foo...')
    
foo()  # Executing foo...
foo()  # <nothing printed as None response was cached and returned>
foo.cache_info()  # CacheInfo(hits=1, misses=1, maxsize=5, currsize=1)
foo.cache_clear()
foo()  # Executing foo... (notice the 'side effect of print is executed again)

Note: debugging this isn’t always obvious. Later on I demonstrate how to mock the builtin open function, and in that scenario I stumbled across this issue, because although I wasn’t mocking the top level function itself (I was mocking the call to open within), the contents of the file being opened was what was returned and being cached.

Mock Module Level/Global Variables

With a module variable you can can either set the value directly or use mock.patch.

In the following example we have the variable client_id which is a global variable inside the app.aws module which we import to reference elsewhere in our code:

import app.aws


def test_account_confirm_successful():
    app.aws.client_id = 456  # used internally by `confirm()`
    ...
    
@mock.patch('app.aws.client_id', 456)
def test_account_confirm_successful():
    ...

In the mock.patch example, there are two key things to notice:

we don’t use return_value.
there is no mock instance passed to the test function.

This is because we’re modifying a variable and not a direct function or ‘callable’ so there’s no need to pass a mock into the test function (if you want to change the value a few times within the test itself then you would mock the variable but not immediately assign a value in the decorator).

Mock Instance Method

There are multiple ways to achieve mocking of an instance method. One common approach is to use mock.patch.object, like so:

from unittest import mock

def test_foo():
    with mock.patch.object(FooClass, 'method_of_class', return_value=None) as mock_method:
        instance = SomeClass()
        instance.method_of_class('arg')
        mock_method.assert_called_with('arg')

Another approach is to mock the method like you would a normal function, but you reference the method via the classname:

def test_bar():
    r = mock.Mock()
    r.content = b'{"success": true}'

    with mock.patch('requests.get', return_value=r) as get:  # Avoid doing actual GET request
        some_function()  # internally calls requests.get
        get.assert_called_once()

Another (although more heavy handed) approach for mocking a class instance method is to take advantage of the fact that a Mock will return a new mock instance when called:

@mock.patch("foo.bar.SomeClass")
def test_stuff(mock_class):
    mock_class.return_value.made_up_function.return_value = "123"

Note: in the above example we mock the entire class, which might not be what you want. If not, then use the previous mock.patch.object example instead.

The reason the above example works is because we’re setting return_value on our mock. Because this is a MagicMock every attribute referenced returns a new mock instance (a function or property you call on a mock doesn’t have to exist) and so we call made_up_function on the returned mock, and on that newly created mock we set the final return_value to 123.

But as mentioned in the above note, this approach might be a little too blunt depending on what your needs are (whether you care whether you have a some what functioning class or not).

Mock Class Method

To mock a class method is a similar approach to mocking an instance method.

One approach might be that you mock the entire class (but now you have one less return_value to assign to):

mock_class.ClassMethodName.return_value = "123"

Or better yet you should mock it like you would any normal function, but just reference the method via the class:

@mock.patch('myapp.Foo.class_method_name')
def test_classmethod(self, mock_class_method):
    mock_class_method.return_value = "foobar"

Mock Entire Class

To mock an entire class you’ll need to set the return_value to be a new instance of the class.

@mock.patch('myapp.app.Car')
def test_class(self, mock_car):

    class NewCar(object):
        def get_make(self):
            return 'Audi'

        @property
        def wheels(self):
            return 6

    mock_car.return_value = NewCar()
    ...

See other class related mocking tips here

Mock Async Calls

Mocking asynchronous code is probably the most confusing aspect of mocking. My ‘go to’ solution I’ll explain first, but after that I’ll share some alternative methods I’ve seen and tried in the past.

First consider this asynchronous code inside of a app.foo module:

import app.stuff

async def do_thing(x):
  return await app.stuff.some_concurrent_function(x)

If we need to mock the coroutine app.stuff.some_concurrent_function, then we can solve this by creating a function that acts as a coroutine and allow it to be configurable for different types of responses:

Note: the example uses tornado for running an asynchronous test.

def make_coroutine(response):
    """You could pass response as a mock or a raw data structure, doesn't matter."""
    
    async def coroutine(*args, **kwargs):
        """*args will be whatever is passed to the original async function.
        
        Meaning you could have a conditional check that let's us change 
        the response to be anything we need.
        """

        return response

    return coroutine

class TestThing(tornado.testing.AsyncTestCase):
    @mock.patch('app.stuff.some_concurrent_function')
    @tornado.testing.gen_test
    def test_async_func(self, mock_thing):
        mock_thing.side_effect = make_coroutine('some response')
        result = yield app.foo.do_thing('xyz')
        assert result == 'some response'

If you do include an if statement within make_coroutine, you could pass in a MagicMock as a simple way of having a single input give you multiple different values back…

def make_coroutine(response):
    async def coroutine(*args, **kwargs):
        if args[0] == 'x':
            return response.x
        elif args[0] == 'y':
            return response.y
        else:
            return response.default

    return coroutine

m = mock.MagicMock(x=1, y=2, default=3)
coro = make_coroutine(m)

When dealing with side_effects that need to sometimes trigger an Exception and other times suceed you could use a slightly modified mock implementation that checks if the given response object is callable or not…

count = 0

def make_side_effect_coroutine(side_effect):
    """Side effect friendly mock coroutine.

    In some tests we need to have a mocked coroutine return a different value
    when it's called multiple times, but a mock side_effect can't trigger a
    raised exception when given an iterator, and so we have to construct that
    behaviour ourselves.
    """

    async def coroutine(*args, **kwargs):
        return side_effect(*args, **kwargs) if callable(side_effect) else side_effect
    return coroutine
    
@mock.patch('app.thing')
def test_confirm_email_change_failure(self, mock_thing):

    def side_effects(*args, **kwargs):
        """Use global var to control mock side effects."""

        global count

        if count > 0:
            raise Exception('whoops')

        count += 1
        return  # don't raise an exception the first time around

    mock_thing.side_effect = make_side_effect_coroutine(side_effects)

If the above approach doesn’t work for you, here are some alternatives…

AsyncMock

Note: this utilizes the package pytest-asyncio to help with testing asyncio code

Let’s start with the code to be mocked…

import asyncio

async def sum(x, y):
    await asyncio.sleep(1)
    return x + y

Now here’s how we’d mock it…

import pytest
import asyncio

# create a new pytest fixture called mock_sum
#
@pytest.fixture()
def mock_sum(mocker):
    async_mock = AsyncMock()
    mocker.patch('app.sum', side_effect=async_mock)
    return async_mock

    # Python <3.8 would have used
    #
    # future = asyncio.Future()
    # mocker.patch('app.sum', return_value=future)
    # return future

@pytest.mark.asyncio
async def test_sum(mock_sum):
    mock_sum.return_value = 4

    # Python <3.8 would have used
    #
    # mock_sum.set_result(4)

    result = await sum(1, 2)
    assert result == 4

Monkey Patch

# allow mock to be used as an await expression...

async def async_response():
    return namedtuple('_', ['body'])('{"state": "success"}')


def mock_async_expression(our_mock):
    return async_response().__await__()


mock.MagicMock.__await__ = mock_async_expression

MagicMock Subclass

class AsyncMock(MagicMock):
    async def __call__(self, *args, **kwargs):
        return super(AsyncMock, self).__call__(*args, **kwargs)
        
class TestHandlers(testing.AsyncTestCase):
    @mock.patch('app.handlers.trigger_soft_cdn_purge', new_callable=AsyncMock)
    @mock.patch('app.handlers.api')
    @testing.gen_test
    async def test_update_cache(self, api_mock, trigger_soft_cdn_purge):
        response = mock.MagicMock()
        response.code = 200
        api_mock.buzz = AsyncMock(return_value=response)

Async Inline Function

@mock.patch('app.buzz_api.api_gateway')
@testing.gen_test
async def test_buzz_api(self, client_mock):
    async def get(url, **kwargs):
        return
        
    client_mock.get.side_effect = get

Mock Instance Types

When mocking an object you’ll find that the mock replaces the entire object and so it can cause tests to pass (or fail) in unexpected ways.

Meaning, if you need to make a mock more like the concrete interface, then there are two ways to do that:

spec
wrap

We can use mock’s spec feature to mimic all methods/attributes of the object being mocked. This ensures your mocks have the same api as the objects they are replacing.

Note: there is a stricter spec_set that will raise an AttributeError.

This is best demonstrated with an example:

import unittest.mock as mock
import tornado.simple_httpclient

from tornado.httpclient import AsyncHTTPClient


http_client = AsyncHTTPClient()
type(http_client)  # tornado.simple_httpclient.SimpleAsyncHTTPClient

isinstance(http_client, tornado.simple_httpclient.SimpleAsyncHTTPClient)  # True

isinstance(mock.MagicMock(), tornado.simple_httpclient.SimpleAsyncHTTPClient)  # False

m = mock.MagicMock(spec=tornado.simple_httpclient.SimpleAsyncHTTPClient)
isinstance(m, tornado.simple_httpclient.SimpleAsyncHTTPClient)  # True

The wrap parameter on the other hand allows you to ‘spy’ on the implementation, as well as affect its behaviour. In the following example I want to spy on the builtin datetime implementation:

@pytest.mark.parametrize("input_date, input_url, valid", [
    ("2017-06-17T00:00:00.000000Z", "foo", True),
    ("2017-06-18T00:00:00.000000Z", "bar", False),
])
@mock.patch("app.handlers.data.datetime", wraps=datetime)
def test_valid_video(mock_datetime, input_date, input_url, valid):
    mock_datetime.now.return_value = datetime(2017, 6, 18, 00, 00, 00, 000000)
    assert valid_video(input_date, input_url) is valid

Another way to use wrap is like so:

from unittest import mock


class Foo(object):
    def bar(self, x, y):
        return x + y + 1

def test_bar():
    foo = Foo()
    with mock.patch.object(foo, 'bar', wraps=foo.bar) as wrapped_foo:
        foo.bar(1, 3)
        wrapped_foo.assert_called_with(1, 2)

test_bar()

Another simplified version of the above that mocks the whole object, not just a single method:

from unittest import mock

class Foo():
    def bar(self, msg):
        print(msg)

f = Foo()
spy = mock.MagicMock(wraps=f)

spy.bar('baz')
spy.bar.assert_called_with('beep')  # raises AssertionError

Mock builtin `open` function

Python’s mock library provides an abstraction for mocking the builtin open function a lot simpler…

def test_load_ui_messages_successful():
    """Verify ui message YAML file can be read properly."""

    file_content = 'foo: bar'

    with mock.patch('bf_auth.utility.open', mock.mock_open(read_data=file_content), create=True) as mock_builtin_open:
        assert utils.load_ui_messages('./path/to/non/existing/file.yaml') == {'foo': 'bar'}

The create=True param set on mock.patch means that the mock.MagicMock returned will automagically create any attributes that are called on the mock (this is because the open function will attempt to access lots of different things and it’s easier for mock to mock out all of that for you).

Conclusion

There we’ll end. Hopefully this list of mocking techniques will be able to see you through even the most complex of code needing to be tested. Let me know what you think on twitter.