UnitTest Python 只模拟一个函数多次调用答案

【问题标题】：UnitTest Python mock only one function multiple callUnitTest Python 只模拟一个函数多次调用
【发布时间】：2013-04-01 20:40:09
【问题描述】：

我正在使用 Mock (http://www.voidspace.org.uk/python/mock/mock.html)，遇到了一个我无法找到解决方案的特定模拟案例。

我有一个函数多次调用 some_function 被 Mocked。

def function():
    some_function(1)
    some_function(2)
    some_function(3)

我只想模拟对 some_function 的第一次和第三次调用。我想第二次调用真正的 some_function。

我用http://www.voidspace.org.uk/python/mock/mock.html#mock.Mock.mock_calls 尝试了一些替代方案，但没有成功。

提前感谢您的帮助。

【问题讨论】：

这能回答你的问题吗？ Mock side effect only X number of times

标签： python unit-testing mocking

【解决方案1】：

看来wraps 参数可能是您想要的：

wraps：要包装的模拟对象的项目。如果 wraps 不是 None 然后调用 Mock 会将调用传递给被包装的对象（返回实际结果并忽略 return_value)。

但是，由于您只想不模拟第二次调用，我建议使用mock.side_effect。

如果 side_effect 是可迭代的，那么每次调用 mock 都会返回可迭代的下一个值。

如果您想为每次调用返回不同的值，这是一个完美的选择：

somefunction_mock.side_effect = [10, None, 10]

只有第一次和第三次调用 somefunction 会返回 10。

但是，如果您确实需要调用 real 函数，但不是第二次，您也可以传递 side_effect 一个可调用对象，但我觉得它很丑陋（可能有一个更聪明的去做）：

 class CustomMock(object):

     calls = 0

     def some_function(self, arg):
         self.calls += 1
         if self.calls != 2:
             return my_real_function(arg)
         else:
             return DEFAULT

somefunction_mock.side_effect = CustomMock().some_function

【讨论】：

【解决方案2】：

比创建CustomMock 类还要简单：

def side_effect(*args, **kwargs):
    if side_effect.counter < 10:
        side_effect.counter += 1
        return my_real_function(arg) 
    else:
        return DEFAULT

side_effect.counter = 0
mocked_method.side_effect = side_effect

【讨论】：

我发现my_real_function 需要别名，就像import my_real_function as someother_name 一样，否则side_effect 函数会一直在无限递归中被调用。

【解决方案3】：

我今天也遇到了同样的情况。经过一番犹豫，我找到了一种不同的解决方法。

起初，我有一个函数是这样的：

def reboot_and_balabala(args):
    os.system('do some prepare here')
    os.system('reboot')
    sys.exit(0)

我希望调用os.system的第一个调用，否则本地文件不会生成，我无法验证它。但我真的不希望调用 os.system 的第二次调用，哈哈。

起初，我有一个类似于以下的单元测试：

def test_reboot_and_balabala(self):
    with patch.object(os, 'system') as mock_system:
        # do some mock setup on mock_system, this is what I looked for
        # but I do not found any easy and clear solution
        with patch.object(sys, 'exit') as mock_exit:
            my_lib.reboot_and_balabala(...)
            # assert mock invoke args
            # check generated files

但最后，我意识到，在调整代码后，我有一个更好的代码结构和单元测试，通过以下方式：

def reboot():
    os.system('reboot')
    sys.exit(0)

def reboot_and_balabala(args):
    os.system('do some prepare here')
    reboot()

然后我们可以通过以下方式测试这些代码：

def test_reboot(self):
    with patch.object(os, 'system') as mock_system:
        with patch.object(sys, 'exit') as mock_exit:
            my_lib.reboot()
            mock_system.assert_called_once_with('reboot')
            mock_exit.assert_called_once_with(0)

def test_reboot_and_balabala(self):
    with patch.object(my_lib, 'reboot') as mock_reboot:
        my_lib.reboot_and_balabala(...)
        # check generated files here
        mock_reboot.assert_called_once()

这不是问题的直接答案。但我认为这是非常鼓舞人心的。

【讨论】：