Flask + SQLAlchemy でテストごとに DB を rollback する

Flask アプリのテストで、簡単なものなら sqlite のメモリDBを使えば良いですが、複雑なアプリだと本番環境と同じ RDBMS を使って実行したいことが多いです。

その際、テストごとにDBを初期化すると遅いので、コミットせずに毎回ロールバックで対応したいのですが、 Flask.test_client や WebTest などを使って複数のHTTPリクエストを行うテストでは、リクエストをまたいでデータを引き継ぐ必要があります。

これを実現するために、

1. テスト中は session.commit() を session.flush(), session.expire_all() に置き換える
2. リクエスト終了時、通常は session.remove() し、テスト中は session.expire_all() する
3. テストごとに session.remove() する。

というカスタマイズをしています。

session.flush() はセッション (Unit of Work) が管理している変更を全てDBに書き出します。
session.expire_all() はセッションが管理しているオブジェクトを全て expire し、次にそのオブジェクトを利用するときはDBから取得し直すようにします。これでメモリ上の値でなく、ちゃんとDBから読んだ値でテストを実行できます。
特に、テスト中のリクエスト終了時に expire_all() することで、 commit() 漏れがあった場合にきちんと次のリクエストでその値が消えています。
session.remove() はそのセッションのトランザクションを rollback し、コネクションをコネクションプールに返します。

カスタマイズ方法を公開しておきます。

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() でエミュレートする."""
        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)
        # テスト時はテスト側でセッションの寿命を管理する.
        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() は app.testing の設定後にする.
    db.init_app(app)

    # ここで view の登録を行う

    return app

test_app.py

import unittest
import myapp

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

        # テストごとに rollback
        self.addCleanup(myapp.db.session.remove)

        @self.app.after_request:
        def after_request(response):
            u"""リクエストごとにコミットしてない変更を忘れる"""
            myapp.db.session.expire_all()
            return response

独自にセッション管理しているのでこのままだと Flask-SQLAlchemy に対応できませんが、 Session.remove() もオーバーライドして通常は expire_all() だけを行うようなカスタマイズをすれば同じことができると思います。

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)