Python mock - mocking class method that modifies class attributes

Question

I currently have the following basic Python class that I want to test:

class Example:

    def run_steps(self):
        self.steps = 0

        while self.steps < 4:
            self.step()
    
    def step(self):
        # some expensive API call
        print("wasting time...")
        time.sleep(1000)

        self.steps += 1

As you can see, the step() method contains an expensive API call so I want to mock it with another function that avoids the expensive API call but still increments self.steps. I found that this is possible by doing this (as seen from here):

def mock_step(self):
    print("skip the wasting time")
    self.steps += 1

# This code works!
def test(mocker):
    example = Example()
    mocker.patch.object(Example, 'step', mock_step)

    example.run_steps()

I simply create a function called mock_step(self) that avoids the API call, and I patch the original slow step() method with the new mock_step(self) function.

However, this causes a new problem. Since the mock_step(self) function is not a Mock object, I can't call any of the Mock methods on it (such as assert_called() and call_count()):

def test(mocker):
    example = Example()
    mocker.patch.object(Example, 'step', mock_step)

    example.run_steps()

    # this line doesn't work
    assert mock_step.call_count == 4

To solve this issue, I have tried to wrap mock_step with a Mock object using the wraps parameter:

def test(mocker):
    example = Example()

    # this doesn't work
    step = mocker.Mock(wraps=mock_step)
    mocker.patch.object(Example, 'step', step)

    example.run_steps()

    assert step.call_count == 4

but then I get a different error saying mock_step() missing 1 required positional argument: 'self'.

So from this stage I am not sure how I can assert that step() has been called exactly 4 times in run_steps().

score 2 · Accepted Answer · answered Jun 02 '21 at 06:37

2

There are several solutions to this, the simplest is probably using a standard mock with a side effect:

def mock_step(self):
    print("skip the wasting time")
    self.steps += 1


def test_step(mocker):
    example = Example()
    mocked = mocker.patch.object(Example, 'step')
    mocked.side_effect = lambda: mock_step(example)
    example.run_steps()
    assert mocked.call_count == 4

side_effect can take a callable, so you can both use a standard mock and the patched method.

answered Jun 02 '21 at 06:37

MrBean Bremen

14,916
3
26
46

I tested your solution with an `import mock` instead of injecting `mocker`. But you seem to patch the class' **definition of the method**, not the actual **method bound to the object**, which results in the time-wasting code being called. – Lenormju Jun 02 '21 at 06:53
@Lenormju - well, I tested the code as is and it works for me (e.g. doesn't call the production code). I used the mocker because the question used it and it is more convenient with `pytest`. Also I wanted to allow using the standard `mock` to allow for `call_count`, `assert_called_xxx` etc. instead of just doing the count myself (which is also a solution, of course). – MrBean Bremen Jun 02 '21 at 08:20
1

you are correct, I missed that it was present in the OP's code, I just had to `pip install pytest-mock` and it worked. Yours is a better solution than mine :) – Lenormju Jun 02 '21 at 08:37
Thanks, this worked for me. I wasn't aware that `mocker.patch` returned a Mock of the patched object so it never occurred to me to do something like `mocked = mocker.patch.object(Example, 'step')`. Thanks for your help :) – runoxinabox Jun 02 '21 at 23:42

score 0 · Answer 2 · answered Jun 02 '21 at 06:48

import unittest.mock as mock
from functools import partial


def fake_step(self):
    print("faked")
    self.steps += 1


def test_api():
    api = Example()
    with mock.patch.object(api, attribute="step", new=partial(fake_step, self=api)):
        # we need to use `partial` to emulate that a real method has its `self` parameter bound at instantiation
        api.run_steps()
    assert api.steps == 4

Correctly outputs "faked" 4 times.

Python mock - mocking class method that modifies class attributes

2 Answers2