Mocking a function returning an object results in AssertionError: Expected '...' to have been called once. Called 0 times

Question

I'm writing a simple example to help me understand how mocking works in unittest. I have a module with two functions:

# model.animals.py
def get_animals(animal_type):
    db = connect_to_db()
    result = db.query_all_data()
    return list(filter(lambda x: x['animal_type'] == animal_type, result))

def connect_to_db():
    pass # That would normally return a DB connection instance

I want to test the get_animals() function which uses a DB connection to retrieve information about all animals and then filters returned data based on animal type. Since I don't want to set up the whole database, I just want to mock the connect_to_db() function which returns a DB connection instance.

This is my test class:

# test_mock.py
from unittest import TestCase, main
from unittest.mock import Mock, patch
from model.animals import get_animals


class GetDataTest(TestCase):

@patch('model.animals.connect_to_db')
def test_get_animals(self, mock_db: Mock):
    mock_db.return_value.query_all_data.return_value = [
        {
            'animal_type': 'meerkat',
            'age': 5
        },
        {
            'animal_type': 'meerkat',
            'age': 11
        },
        {
            'animal_type': 'cow',
            'age': 3
        }
    ]
    result = get_animals('meerkat') # Run the function under test
    mock_db.assert_called_once() # OK
    mock_db.query_all_data.assert_called_once() # AssertionError
    self.assertEqual(len(result), 2) # OK
    self.assertEqual(result[0]['age'], 5) # OK

    if __name__ == "__main__":
        main()

As part of the test I wanted to not only check the filtering of animals based on their type but also whether all the methods inside get_animals() are called.

The test generally works as expected but I get an error when checking whether the query_all_data() function has been called:

AssertionError: Expected 'query_all_data' to have been called once. Called 0 times.

When I add spec=True to my patch I get another error:

AttributeError: Mock object has no attribute 'query_all_data'

Clearly, the function query_all_data is not visible inside the mock even though I set its return value in the test with mock_db.return_value.query_all_data.return_value = ....

What am I missing?

Answer 1

The reason that mock_db.query_all_data.assert_called_once() failed is that it should be mock_db.return_value.query_all_data.assert_called_once().

I have created a helper library to help me generate asserts for mocks so that I won't stumble on such issues as often.

To use it do: pip install mock-generator

Then, in your test place these lines after result = get_animals('meerkat'):

from mock_autogen import generate_asserts
generate_asserts(mock_db)

When you run the test, it would generate the asserts for you (printed to the console and copied to the clipboard):

assert 1 == mock_db.call_count
mock_db.assert_called_once_with()
mock_db.return_value.query_all_data.assert_called_once_with()

You can then edit the generated asserts and use whichever fits your test.