pytest中如何实现自定义装饰器

在 pytest 中，你可以通过编写自定义装饰器来扩展测试函数的功能。这些装饰器可以用于设置测试前后的操作、修改测试行为或添加额外的标记等。下面将详细介绍如何在 pytest 中实现和使用自定义装饰器。

1. 基本概念

• 装饰器（Decorator）：用于修改函数或方法的行为。
• 钩子（Hooks）：pytest 提供了许多内置的钩子函数，允许你在测试生命周期的不同阶段插入自定义逻辑。

2. 实现自定义装饰器

方法一：使用函数装饰器

你可以创建一个普通的 Python 装饰器函数，然后在测试函数上应用它。

# conftest.py 或 test_module.py

import pytest
import functools

def my_decorator(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        print("装饰器前操作")
        result = func(*args, **kwargs)
        print("装饰器后操作")
        return result
    return wrapper

@pytest.mark.parametrize("test_input,expected", [
    ("3+5", 8),
    ("2+4", 6),
    ("6*9", 54),
])
@my_decorator
def test_eval(test_input, expected):
    assert eval(test_input) == expected

解释：

• my_decorator 是一个装饰器函数，它在被装饰的测试函数执行前后打印信息。
• 使用 @pytest.mark.parametrize 来参数化测试用例。
• 装饰器应用于 test_eval 函数，使得每次测试执行时都会经过装饰器的处理。

方法二：使用 pytest 的钩子函数

pytest 提供了许多内置的钩子函数，你可以在这些钩子中插入自定义逻辑。例如，使用 pytest_runtest_setup 和 pytest_runtest_teardown 来在每个测试前后执行操作。

# conftest.py

import pytest

@pytest.fixture(autouse=True)
def my_fixture():
    print("测试前操作")
    yield
    print("测试后操作")

def test_example():
    assert 1 + 1 == 2

解释：

• 使用 @pytest.fixture 定义一个 fixture，并设置 autouse=True 使其自动应用于所有测试函数。
• 在 yield 前后的代码分别在测试前和测试后执行。

3. 高级用法

a. 添加自定义标记

你可以创建装饰器来为测试函数添加自定义标记，然后在 pytest 中根据这些标记执行特定操作。

# conftest.py

import pytest

def slow_test(func):
    func.slow = True
    return func

@pytest.mark.parametrize("test_input,expected", [
    ("3+5", 8),
    ("2+4", 6),
    ("6*9", 54),
])
@slow_test
def test_eval(test_input, expected):
    # 模拟耗时操作
    import time
    time.sleep(2)
    assert eval(test_input) == expected

def pytest_collection_modifyitems(config, items):
    slow_tests = [item for item in items if hasattr(item.function, 'slow')]
    items.sort(key=lambda x: getattr(x.function, 'slow', False))

解释：

• slow_test 装饰器为测试函数添加了一个 slow 属性。
• 在 pytest_collection_modifyitems 钩子中，根据 slow 属性对测试用例进行排序，确保慢速测试最后执行。

b. 参数化装饰器

你可以创建接受参数的装饰器，以增加灵活性。

# conftest.py

import pytest
import functools

def repeat(num_times):
    def decorator(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            for _ in range(num_times):
                result = func(*args, **kwargs)
            return result
        return wrapper
    return decorator

@pytest.mark.parametrize("test_input,expected", [
    ("3+5", 8),
    ("2+4", 6),
    ("6*9", 54),
])
@repeat(num_times=2)
def test_eval(test_input, expected):
    assert eval(test_input) == expected

解释：

• repeat 是一个接受参数的装饰器工厂，返回实际的装饰器。
• @repeat(num_times=2) 将 test_eval 函数执行两次。

4. 注意事项

• 保持装饰器的幂等性：确保装饰器多次应用不会影响测试结果。
• 避免副作用：装饰器应尽量减少对被装饰函数的副作用，除非这是预期的行为。
• 兼容性：确保自定义装饰器与 pytest 的其他功能和插件兼容。

5. 示例项目结构

project/
├── conftest.py
├── test_module.py
└── utils.py

• conftest.py：定义全局的 fixtures 和钩子。
• test_module.py：编写测试函数并应用自定义装饰器。
• utils.py（可选）：存放通用的工具函数或装饰器。

6. 总结

通过自定义装饰器，你可以在 pytest 中灵活地扩展测试功能，满足各种复杂的测试需求。结合 pytest 提供的丰富功能和钩子机制，可以实现高效且可维护的测试套件。

如果你有具体的需求或遇到问题，欢迎进一步提问！