如何编写request_user上下文管理器封装Python测试的patch逻辑？

这问题我太熟了！咱们可以用Python的上下文管理器来封装重复的patch逻辑，既能让测试代码更简洁，还能自动帮你清理补丁，完全符合你想要的with request_user(user):用法。

我推荐两种实现方式，其中用contextlib.contextmanager的写法最简洁Pythonic，先看这个：

方法一：用contextlib.contextmanager装饰器实现 #

from contextlib import contextmanager
from unittest.mock import patch
# 记得替换成你实际的模块路径
from your_app_module import MyAuthenticationClass

@contextmanager
def request_user(user):
    # 封装原本的patch逻辑，利用patch自身的上下文管理器自动处理补丁的应用和移除
    with patch.object(MyAuthenticationClass, 'authenticate', return_value=(user, 'token')) as mock_auth:
        # 如果你的测试代码需要在with块里操作这个mock对象，可以把它yield出去
        yield mock_auth
        # 不需要的话直接写yield就行

使用示例 #

# 先创建你的用户实例
test_user = User(id=1, username="test_user")

with request_user(test_user):
    # 在这里执行需要认证的操作，比如调用接口、测试业务逻辑
    # 此时MyAuthenticationClass.authenticate会被自动补丁，返回你传入的test_user和'token'
    response = client.get("/protected-endpoint")
    assert response.status_code == 200

这个写法的好处是：

• 完全复用了unittest.mock.patch自带的上下文管理器，不需要手动处理补丁的启动和停止
• 代码简洁，一眼就能看懂封装的逻辑
• 如果需要在测试中验证mock的调用情况（比如检查authenticate被调用了几次），可以直接接收yield出来的mock对象：

with request_user(test_user) as mock_auth:
    # 执行操作
    mock_auth.assert_called_once()

方法二：用类实现上下文管理器（更底层的方式） #

如果你想更清楚地控制上下文的每一步，也可以用类来实现__enter__和__exit__方法：

from unittest.mock import patch
from your_app_module import MyAuthenticationClass

class request_user:
    def __init__(self, user):
        self.user = user
        self._patch = None

    def __enter__(self):
        # 创建补丁并启动
        self._patch = patch.object(MyAuthenticationClass, 'authenticate', return_value=(self.user, 'token'))
        mock_auth = self._patch.start()
        return mock_auth

    def __exit__(self, exc_type, exc_val, exc_tb):
        # 不管with块里有没有抛出异常，都停止补丁
        if self._patch:
            self._patch.stop()

使用方式和第一种完全一样，这种写法适合需要在上下文进入/退出时做额外操作的场景（比如记录日志、清理其他资源）。

两种方式都能满足你的需求，推荐第一种，代码更简洁易维护～

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