如何在Python项目单元测试中禁用指定装饰器？

运行单元测试时想禁用选定的装饰器，但当前代码里的disable_decorator_for未实现正确逻辑，导致测试无法捕获预期的ValueError，测试失败。

问题原因 #

被@handle_value_error装饰后的func_that_raise_value_error实际是装饰器返回的wrapper函数，它会捕获ValueError并打印日志，不会向外抛出异常，所以测试里的assertRaises断言自然会失败。

解决方案 #

利用functools.wraps给wrapper函数自动添加的__wrapped__属性（该属性指向被装饰的原函数），我们可以临时把被装饰的函数替换回原函数，测试完成后再恢复，避免影响其他用例。

方案1：实现disable_decorator_for并手动恢复

先修正disable_decorator_for的实现，同时注意测试后的恢复逻辑：

import unittest
import functools
import sys


def handle_value_error(func):
    @functools.wraps(func)
    def wrapper_handle_value_error():
        try:
            return func()
        except ValueError as e:
            print(f"During calling {func.__name__!r} an error was raised: {str(e)}")
    return wrapper_handle_value_error


@handle_value_error
def func_that_raise_value_error():
    raise ValueError


def disable_decorator_for(func, decorator):
    # 获取函数所在的模块
    module = sys.modules[func.__module__]
    # 遍历包装链，找到最原始的函数并替换
    current_func = func
    while hasattr(current_func, '__wrapped__'):
        original_func = current_func.__wrapped__
        setattr(module, func.__name__, original_func)
        current_func = original_func


class TestFunc(unittest.TestCase):
    def test_func_that_raise_value_error(self):
        # 保存原装饰后的函数，用于测试后恢复
        original_wrapped_func = func_that_raise_value_error
        try:
            disable_decorator_for(func=func_that_raise_value_error, decorator=handle_value_error)
            with self.assertRaises(ValueError):
                func_that_raise_value_error()
        finally:
            # 恢复原函数，避免影响其他测试用例
            module = sys.modules[func_that_raise_value_error.__module__]
            setattr(module, func_that_raise_value_error.__name__, original_wrapped_func)

方案2：用unittest.mock.patch临时替换（更简洁）

不需要额外写disable_decorator_for，直接用patch临时替换函数为原函数，测试结束后自动恢复：

import unittest
import functools
from unittest.mock import patch


def handle_value_error(func):
    @functools.wraps(func)
    def wrapper_handle_value_error():
        try:
            return func()
        except ValueError as e:
            print(f"During calling {func.__name__!r} an error was raised: {str(e)}")
    return wrapper_handle_value_error


@handle_value_error
def func_that_raise_value_error():
    raise ValueError


class TestFunc(unittest.TestCase):
    def test_func_that_raise_value_error(self):
        # 获取被装饰的原函数
        original_func = func_that_raise_value_error.__wrapped__
        # 临时替换模块中的函数为原函数
        with patch(f'{__name__}.func_that_raise_value_error', original_func):
            with self.assertRaises(ValueError):
                func_that_raise_value_error()

注意点 #

• 原代码中的测试方法里有笔误：func_thar_raise_value_error应该改为func_that_raise_value_error，否则会报未定义错误。
• 测试后一定要恢复原函数，不然其他测试用例调用该函数时会使用未装饰的版本，导致逻辑错误。

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