Surrender Python Mocking! I Have You Now.

In my previous article I used simple examples to delve into the nuances of mocking in Python. Now I want to mock some more complicated ones.

There are two different aspects that can make mocking tricky:

• What you mock
• What the thing you mock normally returns

“What you mock” can be tricky if you have a line of code that chains many structures together, for instance:

something = AClass.a_class_method()["a_key"].a_method().an_attribute

With regards to “What the thing you mock normally returns”, the more complicated the structure of the thing that is returned, the more complex the mocking will be.

Say you mock a function, that function could be returning any of the built in types. Some of the common ones being str, int, dict, list, function, class, object, iterable and generator.

Or any ridiculous combination of those things, i.e. a function that returns a list of objects, where each object is an iterable that when iterated over returns a generator!

With that in mind, I am going to mock two scenarios:

• a database query that returns an iterable, as this is one keeps cropping up.
• a contrived, overly-complex piece of code of my own making, just to explore how hard it can be.

Mocking a Database Query that Returns an Iterable

If you have used Flask-SQLAlchemy before it is likely you have seen a class like this:

class Animal(db.Model):
    __tablename__ = "animal"
    id = db.Column(db.Integer, primary_key=True)
    type = db.Column(db.String)
    noise = db.Column(db.String)

This class represents the database table “animal” that has three columns: id, type and noise. To use this class to interrogate the database you do it like so:

def get_animal_noises():
    animal_list = Animal.query.all()
    noises = [animal.noise for animal in animal_list]
    return noises

The code:

Animal.query.all()

will generate SQL and run it against whatever database it is configured to hit. I am not expecting you to know or learn Flask or SQLAlchemy, however, this is the code we am going to mock so let’s breakdown what it is we are mocking. The Animal class inherits the query attribute from db.Model and that attribute holds an instance of the BaseQuery class. The BaseQuery class has an all function that returns the result as a list of Animal objects.

The get_animal_noises function uses the Animal class to get all of the animals from the database and build a list of their noises and return that list.

Let us assume that the Animal class and the get_animal_noises live in a module called animals.py, also let’s assume that there is one row on the animal table that contains:

id: 1
type: honey badger
noise: RAAAHH

I could write a function to use this and assert that the noise returned from the database is RAAAHH like so:

def use_get_noise(self):
    noises = get_animal_noises()
    assert(noises[0] == 'RAAAHH')

However, this is a unit test and so it will not have access to the database, therefore let’s write a test that mocks the Animal class:

@mock.patch("animals.Animal")
def test_get_noise(self, mock_animal):
     noises = get_animal_noises()
     assert(noises[0] == 'RAAAHH')

This test fails with a

IndexError: list index out of range 

which wasn’t the error I was expecting. Let’s delve into why, starting with this line:

animal_list = Animal.query.all()

When the Animal class is not mocked Animal.query.all() returns a list of Animal objects, however, when it is mocked it returns a MagicMock object. Therefore in our test, animal_list will be a MagicMock object.

Next we iterate over the animal_list:

noises = [animal.noise for animal in animal_list]

Since animal_list is a MagicMock object, the question is: what happens when you iterate over a MagicMock object (I was expecting it to error with a TypeError: ‘MagicMock’ object is not iterable)?

If we try it in the interpreter:

>>> from unittest.mock import MagicMock
>>> animal_list = MagicMock()
>>> what_is_this = iter(animal_list)
>>> type(what_is_this)
<class 'list_iterator'>

It turns out that a MagicMock object is an iterable by default! When we try to get the first item from the iterator:

>>> next(what_is_this)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
StopIteration

So, a MagicMock object is an iterable that contains no items. Therefore when this line is executed

noises = [animal.noise for animal in animal_list]

noises ends up as an empty list. Hence the IndexError when noises[0] is called.

To make this test work we need to change Animal.query.all() to return a list of Animal objects. Since we are mocking Animal we can do this like so:

@mock.patch("animals.Animal")
def test_get_noise(self, mock_animal):
    # Create test animal
    test_animal = Animal()
    test_animal.noise = "RAAAHH"
    test_animal.type = "honey badger"
    # Set return_value of all to list containing test animal
    mock_animal.query.all.return_value = [test_animal]
    noises = get_animal_noises()
    assert(noises[0] == 'RAAAHH')

So mock_animal is the MagicMock object for the Animal class. We don’t need to set the return_value of mock_animal as the Animal class is not called like a function. The same goes for query, we don’t set the return_value as it is not called like a function. However, all is called hence the line becomes:

mock_animal.query.all.return_value = [test_animal]

Hooray! We have tested our function that makes a database call. WOOOO!

It is worth pointing out that when testing functions that call databases I would write unit tests to test the code that uses the result of the database call, however, I would write integration tests to test the actual call to the database.

Mocking Something Ridiculous

Time for some silliness. Imagine this scenario:

Dragons spawn badgers. When badgers dance, a flower grows. When flowers grow, stars fall. When stars fall, hell breaks loose.

The code for this is:

from contextlib import contextmanager
class Dragon(object):
    @classmethod
    def spawn_badger(cls):
        new_badger = Badger()
        return new_badger
class Badger(object):
    def dance(self):
        print("The badger is dancing!")
        return [Flower()]
class Flower(object):
    def __init__(self):
        print("The flower is growing")
        self.star = {"the_star": Star()}
class Star(object):
    def fall(self):
        print("The stars are falling!")
        yield Hell()
class Hell(object):
    def break_loose(self):
        return True
    @contextmanager
    def initiate_breaking_loose(self):
        print("AHHH! Hell may break loose!")
        yield self.break_loose()
        print("Did hell break loose?")

The function to use this is:

Okay, this rather random function contains a complex combination of structures, as well as the variable hells that is a list Hell objects which has a context manager. The challenge is to change the value of broke_loose to False through mocking so that hell does not break loose. The constraint is that you can only patch the Dragon class. Let’s give this a try.

Assuming all these classes and the break_hell_loose function are in a module called silly.py, I will start by writing a test with no mocking:

def test_if_hell_broke_loose(self):
    broke_loose = break_hell_loose()
    assert(broke_loose == "HELL DID NOT BREAK LOOSE")

This test fails with:

AssertionError: assert 'HELL BROKE LOOSE' == 'HELL DID NOT BREAK LOOSE'

Now I will patch the Dragon class:

@mock.patch('silly.Dragon')
def test_if_hell_broke_loose(self, mock_dragon):
     broke_loose = break_hell_loose()
     assert(broke_loose == "HELL DID NOT BREAK LOOSE")

This still fails. We are going to need to break down the line

Dragon.spawn_badger().dance()[0].star["the_star"].fall()

in some detail. The Dragon class is not called like a function, but spawn_badger and dance are, which means we must set their return_value:

dance_return_mock = mock_dragon.spawn_badger.return_value.dance.return_value

Next the return from dance is accessed like a list i.e. dance()[0] , which calls the python magic method __getitem__, so let’s set the return value of that (this is the mock of the flower object):

flower_mock = dance_return_mock.__getitem__.return_value

Next the star attribute is access but not called like a function, so we don’t need to set return_value:

star_dict_mock = flower_mock.star

Next the star attribute is accessed like a dictionary, which also calls the python magic method __getitem__, this is the mock of the Star object:

star_mock = star_dict_mock.__getitem__.return_value

Next the fall method is called on the Star object, so we need to set its return_value but it needs to return a generator which yields a Hell object. A generator itself is just an iterable, so we can use a simple list instead, and we can use a MagicMock object instead of a Hell object.

hell_mock = mock.MagicMock()    
star_mock.fall.return_value = [hell_mock]

Okay, so far so good. Now the hells variable should contain a list with one item, which is our hell_mock object. However, we now need to look to see how hells is used. The next line in the break_hell_loose function is:

for hell in hells:

We are iterating over hells, and since it is a list of one it means that hell will be our hell_mock object. Let’s now look at how hell is used:

with hell.initiate_breaking_loose() as broke_loose:
    if broke_loose:

The method initiate_breaking_loose is called on it which means we need to set its return_value, however, the return value from initiate_breaking_loose is used as a context manager (the with statement). This means that the python magic method __enter__ is called and the return from that is set to the broke_loose variable. So if we set the return_value of the __enter__ method we will be setting broke_loose:

hell_mock.initiate_breaking_loose.return_value.__enter__.return_value = False

All together the test looks like this:

When this test is run it passes! We have successfully done some insane mocking and prevented hell from breaking loose! I know you would not normally go to these extreme lengths to mock an entire line in this fashion, and it would definitely be far, far easier just to patch the Hell class, but really I just wanted to show you step by step what nasty things you can mock if you need to.

That brings this articles to an end and I hope you enjoyed it. Please recommend it if you did. Also, I would love to hear about any things you found or are finding difficult to mock, so leave me a response and we can chat about it.