Zsolt Pentek
Zsolt Pentek

Reputation: 25

mocking functions from base class

I have a base class where I connect to a mongodb collection and read a date value from an hdfs file, then delete values that are older compared to the value read from this hdfs file. I have some problems with unittesting, so:

In the child class I have the following:

def remove_old_sources(self):
    input_date = self.get_date_from_input()
    mongo_collection = self.connect_to_collection(db, collection)
    deleted_files_count = self.delete_stale_documents(mongo_collection, input_date)

all the functions that are called above comes from the base class. My question is how can I mock them in the following way:

the input_date with a datetime object and the mongo_collection with Collection that comes from mongomock such as:

mocked_collection = mongomock.MongoClient().db.collection
mocked_collection.insert_many([
    {'_id': '1', 'time': datetime(2010, 10, 10, 10, 10, 10)},
    {'_id': '2', 'time': datetime(2015, 10, 10, 10, 10, 10)},
    {'_id': '3', 'time': datetime(2020, 10, 10, 10, 10, 10)},
    {'_id': '4', 'time': datetime(2030, 10, 10, 10, 10, 10)},
    {'_id': '5', 'time': datetime(2040, 10, 10, 10, 10, 10)},
])

This mongomock created a local mongodb in memory if I understood it correctly.

I tried:

@mock.patch("path_to_base_class.get_date_from_input")
@mock.patch("path_to_base_class.connect_to_collection")
def test_remove_old_docs(connect_mock, date_mock)

and then initialize the child class and run remove_old_sources but the mock objects are MagicMock objects like this, I'm not sure how to overwrite them with what I mentioned above. Also tried with a context manager with mock.patch(BaseClass, "function_name", mocked_stuff) and then initialize the Child class but it doesn't work like this.

Any help on this is much appreciated!

Upvotes: 0

Views: 216

Answers (1)

Martijn Pieters
Martijn Pieters

Reputation: 1121486

You get MagicMock objects because that's the standard object type produced when you call a mocked object. You can always replace their return_value attribute if you want something else passed on.

However, you don't need to mock your mongo collection, your code never touches any of that data anyway.

Your code makes 3 calls:

    input_date = self.get_date_from_input()
    mongo_collection = self.connect_to_collection(db, collection)
    deleted_files_count = self.delete_stale_documents(mongo_collection, input_date)

From the first two calls, the code sets two local variables, input_date and mongo_collection, that are then passed to a 3rd method, self.delete_stale_documents().

If your goal is to verify that the code correctly passes those two arguments on to the 3rd method, then all you have to do is mock out the 3 methods and verify that the return value of the first 2 is passed to the 3rd:

@mock.patch("path_to_base_class.delete_stale_documents")
@mock.patch("path_to_base_class.connect_to_collection")
@mock.patch("path_to_base_class.get_date_from_input")
def test_remove_old_docs(date_mock, connect_mock, delete_mock):
    input_date_mock = date_mock.return_value  # a specific MagicMock object
    collection_mock = connect_mock.return_value  # another specific MagicMock object

    # call code-under-test here

    delete_mock.assert_called_with(collection_mock, input_date_mock)

The delete_mock.assert_called_with() would pass only if the result of the connect_mock() and date_mock() calls were used to call delete_mock().

If your test is more involved, with, say, the return value of self.delete_stale_documents(), you can set a specific value as the return value for that call by setting .return_value before the test:

    delete_mock.return_value = 42

Now, deleted_files_count will be set to 42, and you can make further assertions if your code handled that value correctly.

Upvotes: 1

Related Questions