python - 一起使用 pytest fixtures 和 peewee transactions

Question

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

我有这样一种情况，我想在测试中使用两个 fixture，但是让两个 fixture 通过 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):
    # ...

这有效，但阅读 documentation for peewee表明这是容易出错的:

If you attempt to nest transactions with peewee using the transaction() context manager, only the outer-most transaction will be used. However if an exception occurs in a nested block, this can lead to unpredictable behavior, so it is strongly recommended that you use atomic().

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

我唯一的其他想法是使用评估固定装置的顺序来决定将事务放在何处 (which seems to be alphabetical)，但这看起来确实很脆弱。

有办法实现吗？还是我的测试设计需要重新考虑？

Answer 1

如果你想回滚在测试中创建的所有事务，你可以有一个 fixture 来处理事务本身并让模型 fixtures 使用它:

@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):
    # ...

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