[PYTHON] Rollback DB for each test with Flask + SQLAlchemy

For testing Flask apps, you can use sqlite's memory DB for simple ones, but for complex apps you often want to run using the same RDBMS as your production environment.

At that time, it is slow to initialize the DB for each test, so I would like to handle it by rollback every time without committing, but in the test that makes multiple HTTP requests using Flask.test_client or WebTest, it straddles the request. You need to take over the data.

To achieve this,

1. Replace session.commit () with session.flush (), session.expire_all () during testing
2. At the end of the request, usually session.remove (), and during testing session.expire_all ()
3. Do session.remove () for each test.

I am customizing.

session.flush () writes all the changes managed by the session (Unit of Work) to the DB. session.expire_all () expires all the objects managed by the session and retrieves them from the DB the next time you use them. Now you can run the test with the values read from the DB, not the values in memory. In particular, by doing expire_all () at the end of the request under test, if there is a commit () omission, the value will be properly erased in the next request. session.remove () rolls back the transaction for that session and returns the connection to the connection pool.

I will publish the customization method.

sessionmanager.py

from sqlalchemy import create_engine
from sqlalchemy.orm import scoped_session, sessionmaker, Session


class TestingSession(Session):
    """Session for testing."""

    def commit(self):
        u"""commit()Flusn(), expire_all()Emulate with."""
        self.flush()
        self.expire_all()


class SessionManager(object):

    def __init__(self, app=None):
        if app is not None:
            self.init_app(app)

    def init_app(self, app):
        self._create_session(app)
        #During testing, the test side manages session life.
        if not app.testing:
            app.teardown_appcontext(self.exit_sessions)

    def _create_session(self, app, testing=False):
        self.session = scoped_session(sessionmaker(
            bind=create_engine(app.config['DATABASE_DSL']),
            class_=TestingSession if testing else Session,
            expire_on_commit=False))

    def _exit_session(self, response_or_exc):
        self.session.remove()
        return response_or_exc

myapp.py

import flask
from . import sessionmanager

db = sessionmanager.SessionManager()

def get_app(testing=None):
    app = flask.Flask(__name__)

    app.config.from_envvar('MYAPP_SETTING')
    if testing is not None:
        app.testing = testing

    # db.init_app()Is app.After setting testing.
    db.init_app(app)

    #Register the view here

    return app

test_app.py

import unittest
import myapp

class AppTestCase(unittest.TestCase):
    def setUp(self):
        self.app = myapp.get_app(testing=True)

        #Rollback for each test
        self.addCleanup(myapp.db.session.remove)

        @self.app.after_request:
        def after_request(response):
            u"""Forget uncommitted changes on every request"""
            myapp.db.session.expire_all()
            return response

Flask-SQLAlchemy cannot be supported as it is because it manages sessions independently, but I think that the same can be done by overriding Session.remove () and customizing it so that normally only expire_all () is performed.

from flask.ext.sqlalchemy import SQLAlchemy as BaseSQLAlchemy, SignallingSession

class TestSession(SignallingSession):
    def commit(self):
        self.flush()
        self.expire_all()

    def remove(self):
        self.expire_all()

    def real_remove(self):
        super(TestSession, self).remove()


class SQLAlchemy(BaseSQLAlchemy):
    def create_session(self, options):
        if self.app.testing:
            return TestSession(**options)
        else:
            return SignallingSession(**options)