ローカルでDBのテストをする際、DBにテストデータが残っていたり変更されたままにしていると、後続のテストに影響が出てしまい、毎回同じテストができなくなってしまいます。
そのため、テスト実行後に毎回データを削除する方法を、トランザクションをロールバックさせることで実現しました。
＊環境
- MacOS
- Python 2.7.0
- Peewee 3.0
＊Peewee とは
Python用のORMです。
SQLのように直感的に使うことができシンプルなので、学習コストが低く使いやすいというメリットがあります。
現時点では、SQLite、MySQL、Postgresql に対応しています。
Peewee公式ドキュメント
＊Peewee でトランザクションをロールバック
テーブルを作成するには、データクラスを定義し、db.create_tables() をします。
テーブル作成後にトランザクションを操作したい場合は、transaction() でトランザクションを開始、commit() で変更を保存、rollback() でロールバックをすることができます。
from peewee import *
db = SqliteDatabase('my_app.db')
class BaseModel(Model):
class Meta:
database = db
class User(BaseModel):
username = CharField(unique=True)
password = CharField()
email = CharField()
class Message(BaseModel):
user = ForeignKeyField(User, backref='messages')
content = TextField()
def test_database():
db.create_tables([User, Message])
with db.transaction() as tx:
piyo = User.create(username='piyo', password='ppp', email='piyo@com')
# コミットして変更を保存
tx.commit()
huga = User.create(username='huga', password='hhh', email='huga@com')
tom = User.create(username='tom', password='ttt', email='tom@com')
print "======= start ========="
for user in User.select():
print user.username
tx.rollback()
print "======= rollback ========="
for user in User.select():
print user.username
if __name__ == '__main__':
test_database()
<実行結果>
======= start =========
piyo
huga
tom
======= rollback =========
piyo
＊テスト実行後にロールバック
下記はまだロールバックしていない状態です。
このままだとテーブルにデータが残り続けてしまいます。
def test_mock():
db.create_tables([User])
piyo = User.create(username='piyo', password='ppp', email='piyo@com')
huga = User.create(username='huga', password='hhh', email='huga@com')
tom = User.create(username='tom', password='ttt', email='tom@com')
for user in User.select():
print user.username
if __name__ == '__main__':
test_mock()
print "======= end ========="
for user in User.select():
print user.username
<実行結果>
piyo
huga
tom
======= end =========
piyo
huga
tom
そこで処理実行後にロールバックするデコレーターを作成しました。
＊@functools.wraps とは
ラッパー関数を定義するときに、関数デコレータとして呼び出す便宜関数です。つまり、関数をデコレータとして呼び出すことができるようになります。
functools 公式ドキュメント
def mock_transaction(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):
with db.transaction() as tx:
try:
print "======= start ========="
func(*args, **kwargs)
except:
print "======= error ========="
finally:
print "======= finish ========="
tx.rollback()
return wrapper
@mock_transaction
def test_mock():
db.create_tables([User])
piyo = User.create(username='piyo', password='ppp', email='piyo@com')
huga = User.create(username='huga', password='hhh', email='huga@com')
tom = User.create(username='tom', password='ttt', email='tom@com')
for user in User.select():
print user.username
if __name__ == '__main__':
test_mock()
try:
for user in User.select():
print user.username
except:
print "No User Table!"
<実行結果>
======= start =========
piyo
huga
tom
======= finish =========
No User Table!
＊所感
デコレータ関数の中でtry/except/finallyを使っていますが、これがないと他のクラスから呼び出したときに挙動がおかしくなり、ロールバックされない事象が発生しました。
デコレータを自作できるのはとても便利だったので、テスト以外にも活用できそうです。
Pythonのテスト環境を構築するのは初めてで、今回の方法しか知りませんが、他の方法で実現できるのであれば知りたい気持ちです。
Sign up here with your email
ConversionConversion EmoticonEmoticon