Using pytest fixtures and peewee transactions together

Question

I'm writing a set of unit tests using pytest for some database models that are implemented using peewee. I would like to use database transactions (the database is a Postgres one, if that is relevant) in order to roll back any database changes after each test.

I have a situation where I would like to use two fixtures in a test, but have both fixtures clean up their database models via the rollback method, like so:

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

This works, but reading the documentation for peewee suggests that this is error prone:

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().

However, atomic() does not provide a rollback() method. It seems that when managing transactions explicitly, the key is to use an outer-most transaction(), and use savepoint() context managers within that transaction. But in my test code above, both fixtures are on the same "level", so to speak, and I don't know where to create the transaction, and where to create the savepoint.

My only other idea is to use the order that the fixtures get evaluated to decide where to put the transaction (which seems to be alphabetical), but this seems very brittle indeed.

Is there a way to achieve this? Or does my test design need re-thinking?

Answer 1

If you want to rollback all transactions created within a test, you could have a fixture which takes care of the transaction itself and make the model fixtures use it:

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

This way all models are created within the same transaction and rolled back at the end of the test.