Rollback Many Transactions between tests in Flask

Question

My tests take a long time to run and I am trying to rollback transactions between tests instead of dropping and creating the tables between tests.

The issues is that in some tests I do multiple commits.

EDIT: How do I rollback transactions between tests so that tests will run faster

Here is the Base class used for testing.

import unittest
from app import create_app
from app.core import db
from test_client import TestClient, TestResponse


class TestBase(unittest.TestCase):
    def setUp(self):
        self.app = create_app('testing')
        self.app_context = self.app.app_context()
        self.app_context.push()
        self.app.response_class = TestResponse
        self.app.test_client_class = TestClient
        db.create_all()

    def tearDown(self):
        db.session.remove()
        db.drop_all()
        db.get_engine(self.app).dispose()
        self.app_context.pop()

Here is my attempt at rolling back transactions.

class TestBase(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.app = create_app('testing')
        cls.app_context = cls.app.app_context()
        cls.app_context.push()
        cls.app.response_class = TestResponse
        cls.app.test_client_class = TestClient

        db.create_all()

    @classmethod
    def tearDown(cls):
        db.session.remove()
        db.drop_all()
        db.get_engine(cls.app).dispose()

    def setUp(self):
        self.app_content = self.app.app_context()
        self.app_content.push()
        db.session.begin(subtransactions=True)

    def tearDown(self):
        db.session.rollback()
        db.session.close()

        self.app_context.pop()

While many people will argue this. You don't really need to test running database commands. Unit tests are for business logic, which you can then create a mock database to avoid problems like this, and you don't have the risk of messing up your database. — CodeLikeBeaker, Oct 10 '14 at 20:48
Do you use in-memory db for testing? If not, that could dramatically speed up tests. — jsnjack, Oct 14 '14 at 16:18

score 6 · Answer 1 · answered Apr 08 '16 at 20:49

This is the code we use to do this. Make sure that __start_transaction gets called in your setup, and __close_transaction in your teardown (with an app context if you're using flask-sqlalchemy). As a further hint, only inherit this code in test cases that hit the database, and seperate the code that checks your database function from the code that checks your business logic, because those will still run WAY faster.

def __start_transaction(self):
    # Create a db session outside of the ORM that we can roll back
    self.connection = db.engine.connect()
    self.trans = self.connection.begin()

    # bind db.session to that connection, and start a nested transaction
    db.session = db.create_scoped_session(options={'bind': self.connection})
    db.session.begin_nested()

    # sets a listener on db.session so that whenever the transaction ends-
    # commit() or rollback() - it restarts the nested transaction
    @event.listens_for(db.session, "after_transaction_end")
    def restart_savepoint(session, transaction):
        if transaction.nested and not transaction._parent.nested:
            session.begin_nested()

    self.__after_transaction_end_listener = restart_savepoint

def __close_transaction(self):
    # Remove listener
    event.remove(db.session, "after_transaction_end", self.__after_transaction_end_listener)

    # Roll back the open transaction and return the db connection to
    # the pool
    db.session.close()

    # The app was holding the db connection even after the session was closed.
    # This caused the db to run out of connections before the tests finished.
    # Disposing of the engine from each created app handles this.
    db.get_engine(self.app).dispose()

    self.trans.rollback()
    self.connection.invalidate()

score 4 · Answer 2 · answered Oct 16 '14 at 01:23

4

You could use Session.begin_nested. As long as all your tests are properly calling commit to close out their sub-transactions, I think you can simply do

session.begin_nested()
run_test(session)
session.rollback()

Which, in my eyes, seems like it should be faster. Probably depends on your database to some extent, however.

answered Oct 16 '14 at 01:23

vgel

3,225
1
21
35

I replaced db.session.begin(subtransactions=True) with db.session.begin_nested() I get "sqlalchemy.exc.InvalidRequestError: This Session's transaction has been rolled back due to a previous exception during flush. To begin a new transaction with this Session, first issue Session.rollback(). Original exception was: (IntegrityError) duplicate key value violates unique constraint "ix_users_email" DETAIL: Key (email)=(john@example.com) already exists." – Siecje Oct 16 '14 at 13:48
@Siecje that error is due to a different problem, most likely you first need to clear your DB before you start your tests. If an exception happens while running the test, it will stop things before the rollback happens. Then next time you run the test, the table will start already partially filled with data causing this error. – Jeff Widman Jan 13 '16 at 18:52

score 1 · Answer 3 · answered Mar 20 '17 at 22:42

If you're using pytest you can create the following fixtures:

@pytest.fixture(scope='session')
def app():
    app = create_app('config.TestingConfig')
    log.info('Initializing Application context.')

    ctx = app.app_context()
    ctx.push()

    yield app
    log.info('Destroying Application context.')
    ctx.pop()

@pytest.fixture(scope='session')
def db():
    log.info('Initializating the database')

    _db.drop_all()
    _db.create_all()

    session = _db.session
    seed_data_if_not_exists(session)
    session.commit()

    yield _db

    log.info('Destroying the database')
    session.rollback()
    #_db.drop_all() #if necessary

@pytest.fixture(scope='function')
def session(app, db):
    log.info("Creating database session")

    session = db.session
    session.begin_nested()

    yield session

    log.info("Rolling back database session")
    session.rollback()

score 0 · Answer 4 · edited May 23 '17 at 10:30

0

While this answer doesn't technically answer your question, you did mention that the reason behind rolling-back tests is because they take a long time to run, so I'd like to offer an alternate solution:

Create your tables when you start running your test suite, and drop them when all your tests are finished. Then make each test's tearDown simply empty the tables instead of dropping them entirely.

I spent a long time trying to figure out how to speed up my tests via roll-back as the original poster asked, and found it very confusing because it involved nested transactions. However, once I tried the above approach, my test suite ran about twice as fast, which was good enough for me.

edited May 23 '17 at 10:30

Community

1
1

answered Jan 13 '16 at 16:26

Atul Varma

303
3
7

3

This doesn't completely reset everything, notably sequences (such as PK ID numbering) aren't reset. – Jeff Widman Jan 13 '16 at 18:50
This also doesn't work with foreign keys. – Darius Mandres Oct 25 '21 at 14:00

Rollback Many Transactions between tests in Flask

4 Answers4

Linked