Mocking a class method and passing self argument to the Mock's side effect [duplicate]

Question

I am trying to patch a single method in an existing class within a unit test. The class to be patched is:

class Example:
    def __init__(self: "Example", id: int) -> None:
        self.id : int = id
        self._loaded : bool = False
        self._data : typing.Union[str,None] = None

    def data(self: "Example") -> str:
        if not self._loaded:
            self.load()
        return self._data

    def load(self: "Example") -> None:
        self._loaded = True
        # some expensive computations
        self._data = f"real_data{self.id}"

So that instead of calling the self.load() a unittest.mock.Mock is called with the mocked_load function (below) as the side effect:

def mocked_load(self: "Example") -> None:
    # mock the side effects of load without the expensive computation.
    self._loaded = True
    self._data = f"test_data{self.id}"

The first attempt was:

@unittest.mock.patch.object(Example, "load", new = mocked_load)
def test_data__patch_new(
    self: "TestExample",
) -> None:
    example1 = Example(id=1)
    example2 = Example(id=2)

    data1_1 = example1.data()
    self.assertEqual(data1_1, "test_data1")

    data2_1 = example2.data()
    self.assertEqual(data2_1, "test_data2")

    data1_2 = example1.data()
    self.assertEqual(data1_2, "test_data1")

    data2_2 = example2.data()
    self.assertEqual(data2_2, "test_data2")

This works but just replaces the Example.load function with the mocked_load function rather than wrapping it in a Mock; so, although it does pass you cannot extend the test to assert how many times the patched Example.load method was called. It is not the solution I am looking for.

The second attempt was:

@unittest.mock.patch.object(Example, "load")
def test_data__patch_side_effect(
    self: "TestExample",
    patched_load: unittest.mock.Mock
) -> None:
    patched_load.side_effect = mocked_load

    example1 = Example(id=1)
    example2 = Example(id=2)

    self.assertEqual(patched_load.call_count, 0)

    data1_1 = example1.data()
    self.assertEqual(data1_1, "test_data1")
    self.assertEqual(patched_load.call_count, 1)

    data2_1 = example2.data()
    self.assertEqual(data2_1, "test_data2")
    self.assertEqual(patched_load.call_count, 2)

    data1_2 = example1.data()
    self.assertEqual(data1_2, "test_data1")
    self.assertEqual(patched_load.call_count, 2)

    data2_2 = example2.data()
    self.assertEqual(data2_2, "test_data2")
    self.assertEqual(patched_load.call_count, 2)

This fails with the exception:

Traceback (most recent call last):
  File "/usr/lib/python3.8/unittest/mock.py", line 1325, in patched
    return func(*newargs, **newkeywargs)
  File "my_file.py", line 65, in test_data__patch_side_effect
    data1_1 = example1.data()
  File "my_file.py", line 13, in data
    self.load()
  File "/usr/lib/python3.8/unittest/mock.py", line 1081, in __call__
    return self._mock_call(*args, **kwargs)
  File "/usr/lib/python3.8/unittest/mock.py", line 1085, in _mock_call
    return self._execute_mock_call(*args, **kwargs)
  File "/usr/lib/python3.8/unittest/mock.py", line 1146, in _execute_mock_call
    result = effect(*args, **kwargs)
TypeError: mocked_load() missing 1 required positional argument: 'self'

The final attempt was:

@unittest.mock.patch.object(Example, "load")
def test_data__patch_multiple_side_effect(
    self: "TestExample",
    patched_load: unittest.mock.Mock
) -> None:
    example1 = Example(id=1)
    example2 = Example(id=2)

    side_effect1 = lambda: mocked_load( example1 )
    side_effect2 = lambda: mocked_load( example2 )

    patched_load.side_effect = side_effect1

    self.assertEqual(patched_load.call_count, 0)

    data1_1 = example1.data()
    self.assertEqual(data1_1, "test_data1")
    self.assertEqual(patched_load.call_count, 1)

    patched_load.side_effect = side_effect2

    data2_1 = example2.data()
    self.assertEqual(data2_1, "test_data2")
    self.assertEqual(patched_load.call_count, 2)

    patched_load.side_effect = side_effect1

    data1_2 = example1.data()
    self.assertEqual(data1_2, "test_data1")
    self.assertEqual(patched_load.call_count, 2)

    patched_load.side_effect = side_effect2

    data2_2 = example2.data()
    self.assertEqual(data2_2, "test_data2")
    self.assertEqual(patched_load.call_count, 2)

This "works" but its very fragile as the self arguments are hardcoded into the lambda functions and the mock's side_effect needs to be swapped to match each call.

The full minimal representative example is:

import typing
import unittest
import unittest.mock

class Example:
    def __init__(self: "Example", id: int) -> None:
        self.id : int = id
        self._loaded : bool = False
        self._data : typing.Union[str,None] = None

    def data(self: "Example") -> str:
        if not self._loaded:
            self.load()
        return self._data

    def load(self: "Example") -> None:
        self._loaded = True
        # some expensive computations
        self._data = f"real_data{self.id}"

def mocked_load(self: "Example") -> None:
    # mock the side effects of load without the expensive computation.
    self._loaded = True
    self._data = f"test_data{self.id}"

class TestExample( unittest.TestCase ):
    @unittest.mock.patch.object(Example, "load", new = mocked_load)
    def test_data__patch_new(
        self: "TestExample",
    ) -> None:
        # This works but just replaces the Example.load function with another
        # rather than wrapping it in a Mock; so you cannot assert how many
        # times the patched method was called.

        example1 = Example(id=1)
        example2 = Example(id=2)

        data1_1 = example1.data()
        self.assertEqual(data1_1, "test_data1")

        data2_1 = example2.data()
        self.assertEqual(data2_1, "test_data2")

        data1_2 = example1.data()
        self.assertEqual(data1_2, "test_data1")

        data2_2 = example2.data()
        self.assertEqual(data2_2, "test_data2")

    @unittest.mock.patch.object(Example, "load")
    def test_data__patch_side_effect(
        self: "TestExample",
        patched_load: unittest.mock.Mock
    ) -> None:
        # This fails as the self argument is not passed to the side_effect
        # function.

        patched_load.side_effect = mocked_load

        example1 = Example(id=1)
        example2 = Example(id=2)

        self.assertEqual(patched_load.call_count, 0)

        data1_1 = example1.data()
        self.assertEqual(data1_1, "test_data1")
        self.assertEqual(patched_load.call_count, 1)

        data2_1 = example2.data()
        self.assertEqual(data2_1, "test_data2")
        self.assertEqual(patched_load.call_count, 2)

        data1_2 = example1.data()
        self.assertEqual(data1_2, "test_data1")
        self.assertEqual(patched_load.call_count, 2)

        data2_2 = example2.data()
        self.assertEqual(data2_2, "test_data2")
        self.assertEqual(patched_load.call_count, 2)

    @unittest.mock.patch.object(Example, "load")
    def test_data__patch_multiple_side_effect(
        self: "TestExample",
        patched_load: unittest.mock.Mock
    ) -> None:
        # This passes but feels (very) wrong as you have to have change the
        # side_effect each time you call the function and relies on hardcoding
        # the class instances being passed as "self".

        example1 = Example(id=1)
        example2 = Example(id=2)

        side_effect1 = lambda: mocked_load( example1 )
        side_effect2 = lambda: mocked_load( example2 )

        patched_load.side_effect = side_effect1

        self.assertEqual(patched_load.call_count, 0)

        data1_1 = example1.data()
        self.assertEqual(data1_1, "test_data1")
        self.assertEqual(patched_load.call_count, 1)

        patched_load.side_effect = side_effect2

        data2_1 = example2.data()
        self.assertEqual(data2_1, "test_data2")
        self.assertEqual(patched_load.call_count, 2)

        patched_load.side_effect = side_effect1

        data1_2 = example1.data()
        self.assertEqual(data1_2, "test_data1")
        self.assertEqual(patched_load.call_count, 2)

        patched_load.side_effect = side_effect2

        data2_2 = example2.data()
        self.assertEqual(data2_2, "test_data2")
        self.assertEqual(patched_load.call_count, 2)


if __name__ == '__main__':
    unittest.main()

How can I "fix" my second attempt (or give an alternate method) to patch the Example.load function with a Mock so that a side_effect function can be called that has side effects on self?

Answer 1

Straight from this question tbh: Using autospec=True when mocking works from me on python3:

# test.py
import unittest
from unittest.mock import patch


class TestStringMethods(unittest.TestCase):
    @patch.object(Example, "load", autospec=True)
    def test_data__patch_side_effect(
        self: "TestExample",
        patched_load: unittest.mock.Mock
    ) -> None:
        patched_load.side_effect = mocked_load

        example1 = Example(id=1)
        example2 = Example(id=2)

        self.assertEqual(patched_load.call_count, 0)

        data1_1 = example1.data()
        self.assertEqual(data1_1, "test_data1")
        self.assertEqual(patched_load.call_count, 1)

        data2_1 = example2.data()
        self.assertEqual(data2_1, "test_data2")
        self.assertEqual(patched_load.call_count, 2)

        data1_2 = example1.data()
        self.assertEqual(data1_2, "test_data1")
        self.assertEqual(patched_load.call_count, 2)

        data2_2 = example2.data()
        self.assertEqual(data2_2, "test_data2")
        self.assertEqual(patched_load.call_count, 2)

Output:

$ python3 -m unittest test
.
----------------------------------------------------------------------
Ran 1 test in 0.001s

OK

From the docs:

A more powerful form of spec is autospec. If you set autospec=True then the mock will be created with a spec from the object being replaced.

Based on the same paragraph there, you can use an object to define the spec that you want (as an alternative). The following also worked for me:

@patch.object(Example, "load", autospec=Example.load)