Flask + SQLAlchemy + pytest でMySQLを使ってテストごとにロールバックする

西田@大阪です。

今回はPythonの軽量WEBフレームワークであるFlaskと、ORMであるSQLAlchemyを使ってpytestのテストを書いてみたいと思います。

データベースを使った機能をテストするときに、テスト間の依存性をなくすためにテスト毎にデータベースのクリアが必要です。

Flaskの場合SQLiteを使った例が公開されているのですが、MySQLを使った例はありません。今回は Rails のテストなどと同じように、テスト毎にロールバックさせる方式で実装したいと思います。

今回試した環境はこちら

• python: 3.6
• Flask: 1.0.2
• Flask-SQLAlchemy: 2.3.2
• SQLAlchemy: 1.2.9
• pytest: 3.6.2

コミット時の挙動を変更する

commit()メソッドの挙動を変えるためにFlask-SQLAlchemyの SignallingSession を継承したクラスを作成します。

from flask_sqlalchemy import SignallingSession

class TestSignallingSession(SignallingSession):
    def commit(self):
        self.flush() # セッションが保持しているデータをすべてデータベースに書き込む
        self.expire_all() # セッションが保持してるデータをクリアしデータベースより読むこむようにする

上記で作成したクラスを利用するため、SQLAlchemyを継承したクラスを作成し、create_sessionメソッドをオーバーライドします。

from flask_sqlalchemy import SQLAlchemy, orm

class TestSQLAlchemy(SQLAlchemy):
    def create_session(self, options):
        return orm.sessionmaker(class_=TestSignallingSession, db=self, **options)

テスト環境時にコミット時の挙動を変更したセッションが使えるように切り替えるようにします。

from flask import Flask
from flask_sqlalchemy import SQLAlchemy

app = Flask(__name__)

if app.testing:
    db = TestSQLAlchemy(app)
else:
    db = SQLAlchemy(app)

# 以降はこの "db" を使って操作データベースを操作します

テスト終了後にロールバックする

conftest.py にテスト毎にロールバックするように後処理を追加します。

import app # flaskのアプリケーションをimport
import pytest

# 毎テスト毎に rollback() メソッドを実行する
@pytest.fixture(scope='function', autouse=True)
def scope_function():
    yield
    app.db.session.rollback()

さいごに

Flaskは機能が足りてないながらも軽量で、決まった構成がないぶん自由に組めるのが魅力に感じました。

本格的なアプリケーションをつくるのには苦労をしそうですが、少、中規模のアプリケーションを作るにはちょうどよさそうです。

参考

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