CODEWITHSUNDEEP
python
Wayne Werner
Wayne Werner

Reputation: 51847

How do I use a module-level context manager for my tests?

Our database connection code is wrapped in a context manager. It's handy in that all you have to do is something like:

with db.transaction() as cursor:
    do_things_here()

And it handles all the nesting & rolling back of transactions on error for you.

However, it's causing some slight problems for me because I want to wrap my entire module. I thought I could do something like:

@pytest.fixture(scope='module')
def db():
    with db.transaction():
        yield db
        db.rollback()

But that didn't work - pytest says I can't yield from a fixture.

We're using postgres & psycopg2 under the hood.

Is there a way to either use a context manager for the whole module, or solve my issue some other way?

I also tried doing something like this:

@pytest.fixture(scope='module')
def db(request):
    cursor = db.transaction()

    def teardown():
        cursor.__exit__(None, None, None)

    request.addfinalizer(teardown)

but the contextmanager decorator didn't like that.

Finally, I even tried 

with db.transaction(force_rollback=True):
    def test_my_things():
         db.execute('CREATE TABLE castle_auuurrrhhgggg;')

but executing that twice failed because the table exists.

However, if I do:

def test_my_things():
    with db.transaction(force_rollback=True) as cursor:
         cursor.execute('CREATE TABLE knights_who_say_ni;')

It works just fine.

Any ideas of how I can roll everything back at the end of my testing?

Upvotes: 3

Views: 367

Answers (1)

tzaman
tzaman

Reputation: 47820

If you use the yieldfixture instead of a regular fixture it should work with your initial attempt.

Upvotes: 1

Related Questions