Python辅助函数中unittest.mock.patch未作用于实例的问题

我在用unittest.mock.patch编写测试，初始代码运行正常，但想把补丁逻辑封装成方法复用，改写后测试失败了——调用的是真实方法而非被补丁的方法。

初始正常代码 #

class TestService:
  def test_patched(self):
    service = Service()
    with patch.object(service, "_send_to_third_party") as patch_send:
      with patch.object(service, "_notify_listeners") as patch_notify:
        service.do_something()
        patch_send.assert_called_once()
        patch_notify.assert_called_once()

改写后失效的代码 #

class TestService:
  def _patched_service(self):
    service = Service()
    with patch.object(service, "_send_to_third_party") as patch_send:
      with patch.object(service, "_notify_listeners") as patch_notify:
        return service, patch_send, patch_notify

  def test_patched(self):
    service, patch_send, patch_notify = self._patched_service()
    service.do_something()
    patch_send.assert_called_once()
    patch_notify.assert_called_once()

为简化问题，以下测试用例中的第四个test_patching_works_with_function执行失败，我想明确问题原因和可行解决方案：

from unittest.mock import patch

class ExampleClass:
    def __init__(self):
        self.value = 0

    def add(self, value):
        self.value += self._add(value)

    def subtract(self, value):
        self.value -= self._subtract(value)

    def what_is_my_value(self):
        return self.value

    def _add(self, value):
        return value

    def _subtract(self, value):
        return value

def patch_me(exa: ExampleClass):
    with patch.object(exa, '_add', return_value=99):
        with patch.object(exa, '_subtract', return_value=66):
            return exa

class TestPatchingWorks:

    def test_unpatched_works(self):
        exa = ExampleClass()
        exa.add(5)
        exa.subtract(2)
        assert exa.what_is_my_value() == 3

    def test_patching_works_with_patch_class(self):
        exa = ExampleClass()
        with patch.object(ExampleClass, '_add', return_value=30):
            with patch.object(ExampleClass, '_subtract', return_value=10):
                assert exa.what_is_my_value() == 0
                exa.add(5)
                assert exa.what_is_my_value() == 30
                exa.subtract(2)
                assert exa.what_is_my_value() == 20

    def test_patching_works_with_patch_instance(self):
        exa = ExampleClass()
        with patch.object(exa, '_add', return_value=40):
            with patch.object(exa, '_subtract', return_value=30):
                assert exa.what_is_my_value() == 0
                exa.add(5)
                assert exa.what_is_my_value() == 40
                exa.subtract(2)
                assert exa.what_is_my_value() == 10

    def test_patching_works_with_function(self):
        exa = ExampleClass()
        exa = patch_me(exa)
        assert exa.what_is_my_value() == 0
        exa.add(5)
        assert exa.what_is_my_value() == 99
        exa.subtract(2)
        assert exa.what_is_my_value() == 33

问题原因 #

patch.object作为上下文管理器，离开with代码块后会自动撤销补丁。在patch_me或_patched_service函数中，你在with块内返回了实例，但函数执行完毕后，with块的作用域结束，补丁已经被还原，后续调用实例方法时使用的是原始的真实实现。

可行解决方案 #

方案1：将封装方法改为上下文管理器（用yield代替return）

把封装补丁的方法改成生成器，作为上下文管理器使用，这样补丁的作用域会延伸到测试代码的with块内：

针对TestService的修改：

class TestService:
  def _patched_service(self):
    service = Service()
    with patch.object(service, "_send_to_third_party") as patch_send:
      with patch.object(service, "_notify_listeners") as patch_notify:
        yield service, patch_send, patch_notify

  def test_patched(self):
    with self._patched_service() as (service, patch_send, patch_notify):
      service.do_something()
      patch_send.assert_called_once()
      patch_notify.assert_called_once()

针对简化例子的修改：

def patch_me(exa: ExampleClass):
    with patch.object(exa, '_add', return_value=99):
        with patch.object(exa, '_subtract', return_value=66):
            yield exa

# 测试用例修改
def test_patching_works_with_function(self):
    exa = ExampleClass()
    with patch_me(exa) as patched_exa:
        assert patched_exa.what_is_my_value() == 0
        patched_exa.add(5)
        assert patched_exa.what_is_my_value() == 99
        patched_exa.subtract(2)
        assert patched_exa.what_is_my_value() == 33

方案2：用装饰器组合复用补丁

如果要在多个测试方法中复用补丁，可以直接组合@patch装饰器，或者自定义组合装饰器：

class TestService:
  @patch.object(Service, "_send_to_third_party")
  @patch.object(Service, "_notify_listeners")
  def test_patched(self, patch_notify, patch_send):
    service = Service()
    service.do_something()
    patch_send.assert_called_once()
    patch_notify.assert_called_once()

# 自定义复用装饰器
def patch_service_methods(func):
    return patch.object(Service, "_notify_listeners")(patch.object(Service, "_send_to_third_party")(func))

class TestService:
  @patch_service_methods
  def test_patched(self, patch_send, patch_notify):
    service = Service()
    service.do_something()
    patch_send.assert_called_once()
    patch_notify.assert_called_once()

  @patch_service_methods
  def test_another_patched(self, patch_send, patch_notify):
    # 其他测试逻辑
    pass

方案3：手动管理补丁生命周期

通过调用补丁对象的start()和stop()方法，在测试类的setUp和tearDown中统一管理补丁：

class TestService:
  def setUp(self):
    self.service = Service()
    self.patch_send = patch.object(self.service, "_send_to_third_party").start()
    self.patch_notify = patch.object(self.service, "_notify_listeners").start()

  def tearDown(self):
    self.patch_send.stop()
    self.patch_notify.stop()

  def test_patched(self):
    self.service.do_something()
    self.patch_send.assert_called_once()
    self.patch_notify.assert_called_once()

这种方法会让所有测试方法共享补丁逻辑，无需重复编写补丁代码。

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