Better Python mock assertions

25 Sept 2022

• python
• testing

Occasionally I find I’ll end up wanting to be very specific about how a mock has been called but having little or no control over the arguments being passed to it – I want to do mock.assert_called_once_with(arg1, arg2), for example, but it’s awkward to get hold of arg2 from within the test itself, say. Sometimes this is an indication that you might be mocking at the wrong level and you should consider mocking higher or lower in the call stack, but other times it’s unavoidable without rewriting a lot of code or changing the style of an existing test suite.

When I encounter this issue it’s usually in the context of an object instance.Maybe the instance has been generated for my test programmatically, so it’s not always easy to access it directly. If it’s a Django model instance it might use data from faker, so I’d end up having to do something like MyModel.objects.first() to get an instance to use for comparison, which can be brittle. If it’s a class that’s owned by a third-party library then I usually don’t want my tests to know about the inner mechanics of how that library instantiates the class in question. In my hypothetical test I don’t really care which instance of a class is being passed to the mock either, only that it is in fact an instance of that particular class.

Usually this is where things start to get ugly and you start indexing deep into the mock’s call_args_list to make assertions about a single argument used in the call:

assert mock.call_count == 1
assert isinstance(mock.call_args_list[0][0][0], MyModel)

Writing this is kind of painful. It will certainly get your test to pass, but it’s not very readable for someone looking at this test in the future, and I for one can never seem to remember how deep to go into call_args_list, how to get kwargs instead of args and so on. What we really want is something like this:

mock.assert_called_once_with(instance_of_my_model)

If you look at mock.py in the standard library you’ll see that assert_called_once_with is ultimately comparing two _Call objects for equality. This object is a subclass of tuple and has some complex equality logic to handle various scenarios, but it essentially boils down to comparing two two-tuples, each containing the args and kwargs of the expected and actual calls respectively. The order of this comparison is crucial – the expected arguments are compared against the actual arguments. This is what will allow us to make more idiomatic assertions in our test.

You can control how an object compares itself with other objects by declaring the __eq__ method, just like unittest.mock._Call does. Django’s base Model class is a good example of this: two model instances are seen as “equal” if they are derived from the same model and have the same primary key. We know that assert_called_once_with is going to loop over a sequence of expected arguments, comparing them one at a time to see if they are equal to the actual arguments that were used, so our expected argument needs to declare how it should be compared for equality against other objects. Something like this:

class InstanceOf:
    def __init__(self, klass):
        self.klass = klass

    def __eq__(self, other):
      return isinstance(other, self.klass)

Now we can write our test assertion in a much more readable manner:

mock.assert_called_once_with(InstanceOf(MyModel))

As is often the way in this line of work, though, I am far from the first person to encounter this problem! If you find yourself hitting this issue regularly you might want to check out the callee package, which provides a smorgasbord of argument matchers for you to choose from.