pytest fixture创建原理及生成器函数foo的测试方法问询

我定义了一个生成器函数foo，代码如下：

def foo():
    resource = setup()
    yield resource
    tidy(resource)

将其作为pytest fixture使用：

@pytest.fixture
def foofix():
    yield from foo()

该fixture运行正常。我想要测试foo函数，编写了测试代码：

def test_foo():
    res_gen = foo()
    res = next(res_gen)
    assert res.is_working

但调用next后，resource会被立即清理，导致无法断言其是否正常工作。想了解pytest是如何在清理前使用resource的，以及我该如何实现相同的测试方式？

pytest的工作机制 #

pytest处理生成器类型的fixture时，会先执行yield之前的代码获取资源，然后把资源注入给测试函数使用。只有当测试函数执行完毕后，pytest才会继续执行生成器中yield之后的清理代码（也就是tidy(resource)）。这就是fixture能正常使用资源而不会提前被清理的核心原因。

测试foo函数的正确方式 #

要测试foo，你需要在完成断言后，手动驱动生成器执行剩余的清理代码，或者用上下文管理器模拟pytest的逻辑：

方法1：手动控制生成器执行流程 #

def test_foo():
    res_gen = foo()
    res = next(res_gen)
    try:
        assert res.is_working  # 在清理前完成断言逻辑
    finally:
        try:
            next(res_gen)  # 触发生成器执行清理代码
        except StopIteration:
            pass  # 生成器执行完毕抛出该异常属于正常情况，无需处理

用try/finally结构可以确保无论断言是否失败，清理代码都会被执行，避免资源泄漏。

方法2：封装成上下文管理器（更优雅） #

把生成器转换成上下文管理器，模拟fixture的使用体验：

from contextlib import contextmanager

@contextmanager
def foo_context():
    gen = foo()
    resource = next(gen)
    try:
        yield resource
    finally:
        try:
            next(gen)
        except StopIteration:
            pass

def test_foo():
    with foo_context() as res:
        assert res.is_working

这种写法和使用fixture的逻辑一致，上下文管理器会自动在代码块结束后执行清理逻辑。

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