Search code examples
pythonpython-unittestpython-mock

How to replace a method with an arbitrary function inside test function (not using patch.object)?


My attempt was to mock the database operation inside a function with a predefined value.

I patched my mongo collection instance's find method and set a list of dictionary as a demo return value (find return more or less in similar data structure). But the problem is, find returns something that has a count() method, which takes no argument, while, The return value I set (a list) also has a count() method, but it takes a parameter and it's purpose is also different as well.

So, my target was to change the count()'s behavior so that it returns the len of the list I have hard coded. (the len of the return value of the find method)

Below is the code:

In some_module.py,

def somefunc():
    items = mongo_collection.find({"some_field": True}).batch_size(50)

    if items.count() > 0:
        counter += 1

In test_some_module.py,

@patch.object(some_module, 'mongo_collection')
def test_some_func(patched_collection):
    patched_collection.find.return_value.batch_size.return_value = \
                                              [{'id': 1}, {'id': 2}]
    patched_collection.find.return_value.batch_size.return_value.count = ?

Solution

  • It's not clear what you want to test.

    If for some reason you want have list-like "response" and it should act as response (i.e., have count method), you should create such object and set it as return value.

    Now you setting [{'id': 1}, {'id': 2}]. As soon as you return this list via mongo_collection.find().batch_size(), result is actually a list, not mock. So, no further stuff like .count = ... available.

    So, there's ways to go with:

    1. Test response body and your count in different tests, patching connector differently
    2. Create better response mock, i.e.

      class Response(list):
          def count(self):
              return len(self)
      ...
      patched_collection.find.return_value.batch_size.return_value = Response([{'id': 1}, {'id': 2}])
      
    3. Create response mock as instance on response from your mock library