如何简化Python测试中重复使用@patch装饰器的代码？

问题描述 #

我有一段使用mock的Python测试代码（简化版）：

from unittest import TestCase
from mock import patch

class TestClass(TestCase):

    @patch("mymodule.function1")
    @patch("mymodule.function2")
    @patch("mymodule.function3")
    def test_case_1(self, function3_mock, function2_mock, function1_mock):
       # test_case_1代码
       ...

    @patch("mymodule.function1")
    @patch("mymodule.function2")
    @patch("mymodule.function3")
    def test_case_2(self, function3_mock, function2_mock, function1_mock):
       # test_case_2代码
       ...

    @patch("mymodule.function1")
    @patch("mymodule.function2")
    @patch("mymodule.function3")
    def test_case_3(self, function3_mock, function2_mock, function1_mock):
       # test_case_3代码
       ...

    ...

希望简化这段代码，实现两个目标：

1. 无需在每个测试用例函数中重复编写三个@patch(...)语句
2. 无需将mock函数作为参数传入测试用例函数，比如只需定义def test_case_1(self)

请问是否可行？能否提供相关建议或思路？

可行方案 #

完全可以实现，以下是几种实用的思路：

1. 类级别批量Patch #

把@patch装饰器直接应用在测试类上，所有测试方法会自动继承这些patch，不用每个方法重复写：

from unittest import TestCase
from mock import patch

@patch("mymodule.function1")
@patch("mymodule.function2")
@patch("mymodule.function3")
class TestClass(TestCase):
    def test_case_1(self):
        # 通过self.xxx_mock访问mock对象，命名规则是被patch对象的最后部分加"_mock"
        self.function1_mock.return_value = "mock_result_1"
        self.function2_mock.assert_not_called()
        # test_case_1代码
        ...

    def test_case_2(self):
        self.function3_mock.side_effect = ValueError("test error")
        # test_case_2代码
        ...

    def test_case_3(self):
        # test_case_3代码
        ...

类级别的patch会自动把mock对象绑定到测试实例上，比如patchmymodule.function1，对应的属性就是self.function1_mock，直接调用即可。

2. setUp/tearDown手动管理Patch #

如果需要更灵活的控制（比如只给部分方法应用patch），可以在setUp()里启动所有patch，tearDown()里停止，同样不用给测试方法传参数：

from unittest import TestCase
from mock import patch

class TestClass(TestCase):
    def setUp(self):
        # 启动每个patch并把mock对象存到实例属性
        self.function1_patch = patch("mymodule.function1")
        self.function1_mock = self.function1_patch.start()
        
        self.function2_patch = patch("mymodule.function2")
        self.function2_mock = self.function2_patch.start()
        
        self.function3_patch = patch("mymodule.function3")
        self.function3_mock = self.function3_patch.start()

    def tearDown(self):
        # 停止所有patch，清理资源
        self.function1_patch.stop()
        self.function2_patch.stop()
        self.function3_patch.stop()

    def test_case_1(self):
        self.function1_mock.return_value = "test_data"
        # test_case_1代码
        ...

    def test_case_2(self):
        # test_case_2代码
        ...

这种方式还能在setUp()里统一配置mock的默认行为，比如给某个mock设置固定返回值，减少测试方法里的重复代码。

3. 自定义复用装饰器 #

如果多个测试类都需要这组patch，可以封装成自定义装饰器，进一步简化代码：

from mock import patch

def patch_my_functions(cls):
    # 嵌套包装三个patch
    return patch("mymodule.function1")(patch("mymodule.function2")(patch("mymodule.function3")(cls)))

# 使用时直接装饰测试类
@patch_my_functions
class TestClass(TestCase):
    def test_case_1(self):
        self.function1_mock.return_value = "mock_value"
        ...

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