How to mock an attribute variable that returns a class instance

So I have this class located in folder/layer/base.py which has something like this in it:

from folder.plugin import load_plugin

class BaseLayer:

  def __init__(self):
    self.tileindex = load_plugin()

I need to add unit tests to already existing functions within that class. My problem is, the function load_plugin() returns an instance of a class located in folder/tileindex/base.py. Because of that, it happens multiple times and in multiple different functions that a line looks like this:

def somefunction(self):
  key = self.tileindex.get_key(...)
  r = self.tileindex.bulk_add(...)
  self.tileindex.add(...)

And I have no clue how to mock that. At first I was mocking load_plugin and returning whatever value so I could assert it afterwards. But now that I’ve seen these functions that use self.tileindex as an instance of another class I don’t know what to do. For example:

def register(self):
        """
        Registers a file into the system
        :returns: `bool` of status result
        """
        items = [item for item in self.items if item['register_status']]
        if len(items) > 1:
            item_bulk = []
            for item in items:
                item_bulk.append(self.layer2dict(item))
            LOGGER.debug('Adding to tileindex (bulk)')
            r = self.tileindex.bulk_add(item_bulk)
            status = r[items[0]['identifier']]

When I mocked load_plugin, the code failed at the last line saying TypeError: 'Mock' object is not subscriptable.

I tried to import the class that is instanciated and mock that directly. But then for some reason I get the error AttributeError: <Group tileindex> does not have the attribute 'base' as soon as I put @patch('folder.tileindex.base').

Is there some way I can mock self.tileindex itself so I can test the rest of the code?

Thanks!

Answer

Make sure that to not use unittest.mock.Mock but instead unittest.mock.MagicMock for the reasons stated here. You can follow this documentation about Mocking Classes (all of this will use MagicMock).

For your case, here are 3 options to mock the object returned by load_plugin(). You can choose what best fits your needs.

• mock_plugin_return_values – Mocks via return_value
• mock_plugin_side_effect – Mocks via side_effect
• mock_plugin_stub – Mocks via stubbing the class

File tree

.
├── folder
│   ├── layer
│   │   └── base.py
│   ├── plugin.py
│   └── tileindex
│       └── base.py
└── tests
    └── test_layer.py

folder/layer/base.py

from folder.plugin import load_plugin

class BaseLayer:

    def __init__(self):
        self.tileindex = load_plugin()

    def somefunction(self):
        a = self.tileindex.add("a")
        print("add:", a)

        key = self.tileindex.get_key("a")
        print("get_key:", key)

        r = self.tileindex.bulk_add([1, 2, 3])
        print("bulk_add:", r)

        status = r['identifier']
        print("status:", status)

        return a, key, r, status

folder/plugin.py

from folder.tileindex.base import SomePlugin


def load_plugin():
    return SomePlugin()

folder/tileindex/base.py

class SomePlugin():
    pass

test/test_layer.py

import pytest

from folder.layer.base import BaseLayer


# Note, this requires <pip install pytest-mock>


@pytest.fixture
def mock_plugin_return_values(mocker):
    mock_cls = mocker.patch("folder.plugin.SomePlugin")
    mock_obj = mock_cls.return_value

    mock_obj.add.return_value = "Anything!"
    mock_obj.get_key.return_value = "Something!"
    mock_obj.bulk_add.return_value = {"identifier": "Nothing!"}


@pytest.fixture
def mock_plugin_side_effect(mocker):
    mock_cls = mocker.patch("folder.plugin.SomePlugin")
    mock_obj = mock_cls.return_value

    mock_obj.add.side_effect = lambda arg: f"Adding {arg} here"
    mock_obj.get_key.side_effect = lambda arg: f"Getting {arg} now"
    mock_obj.bulk_add.side_effect = lambda arg: {"identifier": f"Adding the {len(arg)} elements"}

@pytest.fixture
def mock_plugin_stub(mocker):
    # Option 1: Create a new class
    # class SomePluginStub:

    # Option 2: Inehrit from the actual class and just override the functions to mock
    from folder.tileindex.base import SomePlugin
    class SomePluginStub(SomePlugin):

        def add(self, arg):
            return f"Adding {arg} here"

        def get_key(self, arg):
            return f"Getting {arg} now"

        def bulk_add(self, arg):
            return {"identifier": f"Adding the {len(arg)} elements"}

    mocker.patch("folder.plugin.SomePlugin", SomePluginStub)


def test_return_values(mock_plugin_return_values):
    layer = BaseLayer()
    result = layer.somefunction()
    print(result)
    assert result == ('Anything!', 'Something!', {'identifier': 'Nothing!'}, 'Nothing!')


def test_side_effect(mock_plugin_side_effect):
    layer = BaseLayer()
    result = layer.somefunction()
    print(result)
    assert result == ('Adding a here', 'Getting a now', {'identifier': 'Adding the 3 elements'}, 'Adding the 3 elements')


def test_stub(mock_plugin_stub):
    layer = BaseLayer()
    result = layer.somefunction()
    print(result)
    assert result == ('Adding a here', 'Getting a now', {'identifier': 'Adding the 3 elements'}, 'Adding the 3 elements')

Output

$ pytest -q -rP
...                                                                                     [100%]
=========================================== PASSES ============================================
_____________________________________ test_return_values ______________________________________
------------------------------------ Captured stdout call -------------------------------------
add: Anything!
get_key: Something!
bulk_add: {'identifier': 'Nothing!'}
status: Nothing!
('Anything!', 'Something!', {'identifier': 'Nothing!'}, 'Nothing!')
______________________________________ test_side_effect _______________________________________
------------------------------------ Captured stdout call -------------------------------------
add: Adding a here
get_key: Getting a now
bulk_add: {'identifier': 'Adding the 3 elements'}
status: Adding the 3 elements
('Adding a here', 'Getting a now', {'identifier': 'Adding the 3 elements'}, 'Adding the 3 elements')
__________________________________________ test_stub __________________________________________
------------------------------------ Captured stdout call -------------------------------------
add: Adding a here
get_key: Getting a now
bulk_add: {'identifier': 'Adding the 3 elements'}
status: Adding the 3 elements
('Adding a here', 'Getting a now', {'identifier': 'Adding the 3 elements'}, 'Adding the 3 elements')
3 passed in 0.06s