ContextDecorator的__exit__未恢复Mock对象，Stripe测试Mock异常求助

嘿，我之前在Mock第三方API时也踩过这个坑！你的问题大概率是在__enter__里没有正确保存原函数的引用，或者__exit__方法里没做好恢复操作。让我拆解一下常见的错误点和修复方案：

常见错误原因 #

• 没有保存原函数的引用：如果是手动替换函数（而非用unittest.mock.patch），你需要在__enter__里把被Mock的原函数存下来，不然__exit__根本不知道要恢复成什么。
• 未正确执行恢复逻辑：如果用了patch对象，你得确保在__exit__里调用每个patch的stop()方法；如果是手动替换，要把存好的原函数重新赋值回去。
• 遗漏部分Mock的函数：如果同时Mock了多个Stripe方法，可能漏恢复其中某一个，导致后续测试受影响。

正确的实现示例 #

下面是一个能正确恢复原函数的MockStripe类，既可以当装饰器也可以当上下文管理器：

import contextlib
from unittest.mock import patch, MagicMock

class MockStripe(contextlib.ContextDecorator):
    def __enter__(self):
        # 定义所有需要Mock的Stripe方法
        self.mock_charge_create = MagicMock(return_value={'id': 'test_charge_123'})
        self.mock_customer_create = MagicMock(return_value={'id': 'test_cust_456'})

        # 创建patch对象并启动，同时保存这些patch以便后续停止
        self.patches = [
            patch('stripe.Charge.create', self.mock_charge_create),
            patch('stripe.Customer.create', self.mock_customer_create)
        ]
        
        for p in self.patches:
            p.start()
        
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # 逐个停止patch，自动恢复原函数
        for p in self.patches:
            p.stop()
        
        # 返回False让异常正常传播（如果测试中有异常的话）
        return False

为什么这个方案有效？ #

• unittest.mock.patch本身已经封装了“保存原函数-替换为Mock-恢复原函数”的逻辑，我们只需要在__enter__启动所有patch，__exit__停止它们就行。
• 不管是用@MockStripe()装饰测试函数，还是用with MockStripe():上下文，__exit__都会被调用，确保所有Mock都被清理。

如果你是手动替换函数（不用patch） #

如果是自己手动替换而不是用patch，一定要记得保存原函数并在__exit__恢复：

# 手动替换的正确写法示例
class MockStripe(contextlib.ContextDecorator):
    def __enter__(self):
        # 保存原函数引用
        self.original_charge_create = stripe.Charge.create
        self.original_customer_create = stripe.Customer.create
        
        # 替换为Mock函数
        stripe.Charge.create = lambda **kwargs: {'id': 'test_charge'}
        stripe.Customer.create = lambda **kwargs: {'id': 'test_customer'}
        
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # 恢复原函数
        stripe.Charge.create = self.original_charge_create
        stripe.Customer.create = self.original_customer_create
        return False

最后检查点 #

• 确认__exit__方法里的恢复逻辑覆盖了所有被Mock的函数；
• 测试时可以在Mock之后和退出上下文/装饰器之后，打印被Mock函数的id()，看是否回到了原来的引用；
• 如果用了patch，不要忘记调用stop()，只调用start()不stop()肯定会导致Mock残留。

内容的提问来源于stack exchange，提问作者Geoffrey R.