Python模拟 - 模拟修改类属性的类方法

Question

我目前有以下基本的Python类需要测试：

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

从上面的代码可以看出，step() 方法包含一个昂贵的 API 调用，因此我希望使用另一个函数来模拟它，该函数避免了昂贵的 API 调用，但仍会增加 self.steps。我发现这是可能的，方法如下（参见这里）：

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()

我创建了一个名为mock_step(self)的函数，它避免了API调用，并将原始较慢的step()方法替换为新的mock_step(self)函数。 然而，这导致了一个新问题。由于mock_step(self)函数不是Mock对象，我无法在其上调用任何Mock方法（如assert_called()和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

为了解决这个问题，我已经尝试使用wraps参数将mock_step包装成一个Mock对象：

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

但接着我收到不同的错误信息，指出：mock_step() missing 1 required positional argument: 'self'。

所以现在我并不确定如何断言run_steps()中step()已被调用了确切的4次。

- runoxinabox

Answer 1

有几种解决方法，最简单的可能是使用带有副作用的标准模拟：

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可以使用可调用对象，因此您既可以使用标准的mock对象，也可以使用已打补丁的方法。

- MrBean Bremen

---

我使用了 import mock 而不是注入 mocker 来测试您的解决方案。但是您似乎修补了类的 方法定义，而不是实际绑定到对象的 方法，这导致了浪费时间的代码被调用。 - Lenormju

@Lenormju - 嗯，我按原样测试了代码，它对我有效（例如，不调用生产代码）。我使用模拟器是因为问题中使用了它，并且在pytest中更方便。此外，我想允许使用标准的mock来允许call_count，assert_called_xxx等，而不仅仅是自己计算数量（当然，这也是一种解决方案）。 - MrBean Bremen

你是正确的，我错过了它在 OP 的代码中已经存在这一事实，我只需要 pip install pytest-mock 就可以了。你的解决方案比我的更好 :) - Lenormju

谢谢，这对我很有帮助。我不知道 mocker.patch 返回的是被修补对象的 Mock，所以我从来没有想过要做像 mocked = mocker.patch.object(Example, 'step') 这样的事情。感谢您的帮助 :) - runoxinabox

Answer 2

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

正确输出"faked" 4次。

- Lenormju