How to rollback database changes during Flask testing

Question

When running tests in pytest, the database is modified. What is the best way to undo changes to the database?

DBSession rollback

For those tests where I can access the backend directly, I currently use pytest fixture to start a new DBSession for every test function, and rollback the session at the end of it

@pytest.fixture(scope='session')
def db(app, request):
    """Session-wide test database."""
    def teardown():
        pass

    _db = SQLAlchemy(app)    
    return _db

@pytest.fixture(scope='function')
def db_session(db, request):
    """Creates a new database session for a test."""
    engine = create_engine(
                            TestConfig.SQLALCHEMY_DATABASE_URI,
                            connect_args={"options": "-c timezone=utc"})
    DbSession = sessionmaker(bind=engine)
    session = DbSession()
    connection = engine.connect()
    transaction = connection.begin()
    options = dict(bind=connection, binds={})
    session = db.create_scoped_session(options=options)
    db.session = session

    yield session

    transaction.rollback()
    connection.close()
    session.remove()

In the test code, I simply use the fixture

def test_create_project(db_session): project = _create_test_project(db_session) assert project.project_id > 0

Flask / HTTP Testing

But for testing the API via Flask/HTTP, I cannot use db_session. Even when I create a fixture to explicitly DROP the test database and restore from production, it will not work because there is no direct database code

@pytest.fixture(scope='function')
def db_session_refresh(db, request):
    """Refresh the test database from production after running the test"""
    engine = create_engine(
                            TestConfig.SQLALCHEMY_DATABASE_URI,
                            connect_args={"options": "-c timezone=utc"})
    DbSession = sessionmaker(bind=engine)
    session = DbSession()

    connection = engine.connect()
    transaction = connection.begin()

    options = dict(bind=connection, binds={})
    session = db.create_scoped_session(options=options)

    db.session = session

    yield session

    transaction.rollback()
    connection.close()
    session.remove()  
    refresh_test_db_sql = """
        SELECT pg_terminate_backend(pid)
        FROM pg_stat_activity
        WHERE datname = 'appdb_test';

        DROP DATABASE appdb_test;

        CREATE DATABASE appdb_test TEMPLATE appdb;
        """
    engine.execute(refresh_test_db_sql)

Even if this works, it is inefficient to refresh the database for every function.

What is the proper/better way to run test that modifies the database?

For sure you should use a separate database for tests instead of rolling back the main one. You set up a database in setUp method and destroy in tearDown. — Aleksandr Borisov
I do have dbapp and dbapp_test. How can I roll back changes to database when running Flask HTTP test code? — Hanxue

akiva akiva · Accepted Answer · 2017-10-29T13:03:59

as commented earlier - creation and destruction of the DB should be taken care out of you unit test and moved into a wrapper.

but to answer you question - I have the impression that you're using postgresql. try removing the database name from your connection string and connect to it directly and add it to the search_path

How to rollback database changes during Flask testing

DBSession rollback

Flask / HTTP Testing

1 Answers