一起使用pytest固定装置和peewee交易

Question

我正在编写一组单元测试，pytest用于一些使用实现的数据库模型peewee。我想使用数据库事务（如果相关，数据库是Postgres数据库），以便在每次测试后回滚任何数据库更改。

我有一种情况，我想在测试中使用两个固定装置，但是要使两个固定装置都通过rollback方法清理其数据库模型，如下所示：

@pytest.fixture
def test_model_a():
    with db.transaction() as txn:  # `db` is my database object
        yield ModelA.create(...)
        txn.rollback()

@pytest.fixture
def test_model_b():
    with db.transaction() as txn:  # `db` is my database object
        yield ModelB.create(...)
        txn.rollback()

def test_models(test_model_a, test_model_b):
    # ...

这行得通，但是阅读文档peewee说明这很容易出错：

如果尝试使用transaction()上下文管理器与peewee嵌套事务，则仅使用最外面的事务。但是，如果嵌套块中发生异常，则可能导致无法预料的行为，因此强烈建议您使用atomic()。

但是，atomic()没有提供一种rollback()方法。似乎在显式管理事务时，关键是使用最外部的transaction()，并savepoint()在该事务内使用上下文管理器。但是可以说，在我上面的测试代码中，两个固定装置都处于同一“级别”上，我不知道在哪里创建事务以及在哪里创建保存点。

我唯一的另一个想法是使用对固定装置进行评估的顺序来决定将交易放在何处（这似乎是字母顺序的），但这确实非常脆弱。

有没有办法做到这一点？还是我的测试设计需要重新考虑？

Answer 1

如果要回滚测试中创建的所有事务，则可以使用一个夹具来处理事务本身，并让模型夹具使用它：

@pytest.fixture
def transaction():
    with db.transaction() as txn:  # `db` is my database object
        yield txn
        txn.rollback()

@pytest.fixture
def test_model_a(txn):
    yield ModelA.create(...)        

@pytest.fixture
def test_model_b(txn):
    yield ModelB.create(...)            

def test_models(test_model_a, test_model_b):
    # ...

这样，所有模型都在同一事务中创建，并在测试结束时回滚。