Python使用unittest.mock.patch.object打补丁后如何恢复被mock方法原始行为

你可以提前保存未被mock的原始方法引用，指定次数的异常抛出完成后直接调用原始方法即可，不需要复制原方法的业务逻辑，修改后的可运行代码如下：

import unittest
from unittest.mock import patch


class Bar(object):

    def some_method_that_may_fail(self):
        return "bar"


class Foo(object):
    bar = None

    def retry_method(self):
        try:
            self.__class__.bar = Bar().some_method_that_may_fail()
        except KeyError:
            self.retry_method()


class TestRetry(unittest.TestCase):

    def setUp(self):
        self.instance = Foo()

    def test_retry(self):
        # 提前保存未被mock的原始方法引用
        original_method = Bar.some_method_that_may_fail
        errors_list = [KeyError("")] * 5

        def raise_errors(*args, **kwargs):
            if errors_list:
                errors_list.pop(0)
                # 前5次调用抛出指定异常
                raise KeyError("")
            # 异常抛完后调用原始方法走原生逻辑
            return original_method(*args, **kwargs)

        with patch.object(Bar, 'some_method_that_may_fail', new=raise_errors) as mocked:
            self.instance.retry_method()
        self.assertEqual(self.instance.bar, 'bar')
        # 可选校验：总共调用6次，前5次失败第6次成功
        self.assertEqual(mocked.call_count, 6)


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

核心实现逻辑 #

• 执行patch.object替换方法前，先将原生的Bar.some_method_that_may_fail赋值给临时变量，此时方法还未被mock替换，拿到的是完整的原始业务逻辑
• 自定义的替换方法中，错误列表不为空时抛出指定异常，列表清空后直接调用提前保存的原始方法，传入调用时的所有参数即可对齐原生调用效果
• 如果不想自己写闭包逻辑，也可以利用mock对象的side_effect特性实现，写法更简洁：

def test_retry(self):
    original_method = Bar.some_method_that_may_fail
    def call_original(*args, **kwargs):
        return original_method(*args, **kwargs)
    with patch.object(Bar, 'some_method_that_may_fail') as mocked:
        # side_effect按列表顺序依次执行，前5次抛异常，最后一次调用原始方法
        mocked.side_effect = [KeyError("")]*5 + [call_original]
        self.instance.retry_method()
    self.assertEqual(self.instance.bar, 'bar')

内容的提问来源于stack exchange，提问作者ronkov