factory_boyで楽々テストデータ生成

Python Advent Calendar 2013 の14日目です。

今日は、テストコードを書くときに便利なモジュール factory_boy を紹介します。

なお、この記事のサンプルコードは以下の環境で動作確認しています。

  • Python 3.3.3(2.7でも動くように書いています)
  • factory_boy 2.2.1
  • SQLAlchemy 0.8.4
  • SQLite3

インストール

インストールは、pipコマンド一発でできます:

pip install factory-boy

今回はSQLAlchemyも使うので、以下のコマンドでインストールしておきましょう:

pip install sqlalchemy

基本的な使い方

以下のコードをjojo.py という名前で保存してください:

import factory
from factory.alchemy import SQLAlchemyModelFactory
from sqlalchemy import Column, Integer, Unicode, create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import scoped_session, sessionmaker

session = scoped_session(sessionmaker())
engine = create_engine('sqlite://')
session.configure(bind=engine)
Base = declarative_base()


class Character(Base):
    u""" 登場人物 """
    __tablename__ = 'character'

    id = Column(Integer(), primary_key=True)
    name = Column(Unicode(20))

    def __str__(self):
        return self.name


Base.metadata.create_all(engine)


class CharacterFactory(SQLAlchemyModelFactory):
    FACTORY_FOR = Character
    FACTORY_SESSION = session   # the SQLAlchemy session object

    id = factory.Sequence(lambda n: n)
    name = factory.Sequence(lambda n: u'登場人物 %d' % n)

jojo.pyと同じディレクトリ上でpythonコマンドを実行して、以下のコードを入力してください:

>>> from jojo import session, Character, CharacterFactory
>>> session.query(Character).all() # ここではまだ何のデータも作られていない
[]
>>> CharacterFactory() # ここでfactory_boyがデータを生成する
<jojo.Character object at 0x10f08f8d0>
>>> session.query(Character).all() # データが1件追加されている
[<jojo.Character object at 0x10f08f8d0>]
>>> print(session.query(Character).all()[0])
登場人物 1

で、何が嬉しいの?

これだけだと何が嬉しいのか分かりませんね。では、コードをもう少し複雑にしてみましょう:

import factory
from factory.alchemy import SQLAlchemyModelFactory
from sqlalchemy import Column, Integer, Unicode, create_engine, ForeignKey
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import scoped_session, sessionmaker, relationship, backref

session = scoped_session(sessionmaker())
engine = create_engine('sqlite://')
session.configure(bind=engine)
Base = declarative_base()


class Stand(Base):
    u"""スタンド"""
    __tablename__ = 'stand'
    id = Column(Integer(), primary_key=True)
    name = Column(Unicode(20))

    def __str__(self):
        return self.name


class Character(Base):
    u""" 登場人物 """
    __tablename__ = 'character'

    id = Column(Integer(), primary_key=True)
    name = Column(Unicode(20))
    stand_id = Column(Integer, ForeignKey('stand.id'))
    stand = relationship('Stand', backref=backref('character', uselist=False))

    def __str__(self):
        return "本体名:{0} スタンド名:{1}".format(self.name, self.stand.name)


Base.metadata.create_all(engine)


class StandFactory(SQLAlchemyModelFactory):
    FACTORY_FOR = Stand
    FACTORY_SESSION = session

    id = factory.Sequence(lambda n: n)
    name = factory.Sequence(lambda n: u'スタンド %d' % n)


class CharacterFactory(SQLAlchemyModelFactory):
    FACTORY_FOR = Character
    FACTORY_SESSION = session   # the SQLAlchemy session object

    id = factory.Sequence(lambda n: n)
    name = factory.Sequence(lambda n: u'登場人物 %d' % n)
    stand = factory.SubFactory(StandFactory)

CharacterモデルとStandモデルが1対1の関係になるように変更しました。再度pythonコマンドを実行して、以下のコードを入力してください:

>>> from jojo import session, Character, CharacterFactory
>>> session.query(Character).all() # ここではまだ何のデータも作られていない
[]
>>> CharacterFactory() # ここでfactory_boyがデータを生成する
<jojo.Character object at 0x105409910>
>>> print(session.query(Character).all()[0]) # Characterに紐づくStandも同時に生成されている
本体名:登場人物 1 スタンド名:スタンド 1

先ほどは、Characterモデルのデータが1個作られただけですが、今度はCharacterに紐づくStandも同時に生成されました。

さて、Character, Standモデルをテストコードの中で使う場面を想像してみてください。テストメソッド実行前にテストデータを投入する際、通常はどのような書き方になるでしょうか?

おそらく、テストデータをfixtureファイル(YAMLとかJSONとかXML)に書いて、それを読み込んでDBに登録するような形になるのではないかと思います。

fixtureファイルにCharacter, Standモデルのようにリレーションが存在するデータを書く際、どうしてもプライマリーキーを意識しなければなりません。

Characterに紐づくStandのidにうっかり実在しないものを書いてしまうと、当然テストデータ投入でコケてしまいます。

今回の例のリレーションは1体1のシンプルなものですが、実際のソフトウェア開発でのリレーションはもっと複雑です。1対多、多対多になると、まず頭のなかがグチャグチャになると思います(というか、私はなります)。

factory_boyを使うと、このfixtureファイル特有の煩わしさから開放されます。

10件テストデータを投入したいなら、こんなコードを書くだけで済みます:

>>> from jojo import session, Character, CharacterFactory
>>> for i in range(10):
...     CharacterFactory()
...
<jojo.Character object at 0x10ac37450>
<jojo.Character object at 0x10ac37e10>
<jojo.Character object at 0x10ac37250>
<jojo.Character object at 0x10ac6e750>
<jojo.Character object at 0x10ac6e810>
<jojo.Character object at 0x10ac6e910>
<jojo.Character object at 0x10ac37190>
<jojo.Character object at 0x10ac37dd0>
<jojo.Character object at 0x10ac6ebd0>
<jojo.Character object at 0x10ac6eb10>
>>> for character in session.query(Character).all():
...     print(character)
...
本体名:登場人物 1 スタンド名:スタンド 1
本体名:登場人物 2 スタンド名:スタンド 2
本体名:登場人物 3 スタンド名:スタンド 3
本体名:登場人物 4 スタンド名:スタンド 4
本体名:登場人物 5 スタンド名:スタンド 5
本体名:登場人物 6 スタンド名:スタンド 6
本体名:登場人物 7 スタンド名:スタンド 7
本体名:登場人物 8 スタンド名:スタンド 8
本体名:登場人物 9 スタンド名:スタンド 9
本体名:登場人物 10 スタンド名:スタンド 10

「登場人物1」とか「スタンド1」とかつまんない。もっと本物っぽいデータがほしいんだ!

データの内容はCharacterFactory, StandFactory に定義した内容に基づいて自動的に生成されますが、こんな風に具体的な値を指定することもできます:

>>> from jojo import session, Character, CharacterFactory, StandFactory
>>> CharacterFactory(name=u'空条承太郎', stand=StandFactory(name='スタープラチナ'))
>>> print(session.query(Character).all()[0])
本体名:空条承太郎 スタンド名:スタープラチナ

また、 Fuzzy attributes を利用することで、ランダムな値を生成することもできます。

例えば、こう書くと:

# Stand, StandFactory以外のコードは省略
import factory
from factory import fuzzy
from factory.alchemy import SQLAlchemyModelFactory
from sqlalchemy import Column, Integer, Unicode


class Stand(Base):
    u"""スタンド"""
    __tablename__ = 'stand'
    id = Column(Integer(), primary_key=True)
    name = Column(Unicode(20))
    power = Column(Integer())
    speed = Column(Integer())
    firing_range = Column(Integer())
    enduring = Column(Integer())
    accuracy = Column(Integer())
    potential = Column(Integer())


class StandFactory(SQLAlchemyModelFactory):
    FACTORY_FOR = Stand
    FACTORY_SESSION = session

    id = factory.Sequence(lambda n: n)
    name = factory.Sequence(lambda n: u'スタンド %d' % n)
    power = fuzzy.FuzzyInteger(1, 5)
    speed = fuzzy.FuzzyInteger(1, 5)
    firing_range = fuzzy.FuzzyInteger(1, 5)
    enduring = fuzzy.FuzzyInteger(1, 5)
    accuracy = fuzzy.FuzzyInteger(1, 5)
    potential = fuzzy.FuzzyInteger(1, 5)

power, speed, firing_range, enduring, accuracy, potential に1から5までのランダムな整数値が生成されます:

>>> from jojo import session, Stand, StandFactory
>>> for i in range(3):
...     StandFactory()
...
<jojo.Stand object at 0x10b025610>
<jojo.Stand object at 0x10b11db90>
<jojo.Stand object at 0x10b11dcd0>
>>> for stand in session.query(Stand).all():
...     print(stand.name, ":", stand.power, stand.speed, stand.firing_range, stand.enduring, stand.accuracy, stand.potential)
...
スタンド 1 : 5 5 2 2 5 1
スタンド 2 : 4 2 1 2 5 5
スタンド 3 : 4 1 5 5 3 4

文字列の場合は:

# Character, Characterfactory以外のコードは省略
import factory
from factory import fuzzy
from factory.alchemy import SQLAlchemyModelFactory
from sqlalchemy import Column, Integer, Unicode, ForeignKey
from sqlalchemy.orm import relationship, backref


class Character(Base):
    u""" 登場人物 """
    __tablename__ = 'character'

    id = Column(Integer(), primary_key=True)
    name = Column(Unicode(20))
    signature_phrase = Column(Unicode(50))
    stand_id = Column(Integer, ForeignKey('stand.id'))
    stand = relationship('Stand', backref=backref('character', uselist=False))

    def __str__(self):
        return "本体名:{0} スタンド名:{1}".format(self.name, self.stand.name)


class CharacterFactory(SQLAlchemyModelFactory):
    FACTORY_FOR = Character
    FACTORY_SESSION = session   # the SQLAlchemy session object

    id = factory.Sequence(lambda n: n)
    name = factory.Sequence(lambda n: u'登場人物 %d' % n)
    signature_phrase = fuzzy.FuzzyChoice(
        choices=(
            u'このド低能がァーーッ',
            u'このクサレ脳ミソがァーーッ',
            u'このきたならしい阿呆がァーッ!!',
            u'この花京院典明に精神的動揺による操作ミスは決してない!と思っていただこうッ!',
        )
    )
    stand = factory.SubFactory(StandFactory)

choicesで指定した文字列のどれかが選択されます:

>>> from jojo import session, Character, CharacterFactory
>>> for i in range(10):
...     CharacterFactory()
...
<jojo.Character object at 0x10ce9e310>
<jojo.Character object at 0x10ce9e990>
<jojo.Character object at 0x10ce9ee50>
<jojo.Character object at 0x10ced2c10>
<jojo.Character object at 0x10ced2c50>
<jojo.Character object at 0x10ced2d90>
<jojo.Character object at 0x10ced2cd0>
<jojo.Character object at 0x10ced2dd0>
<jojo.Character object at 0x10ced8050>
<jojo.Character object at 0x10ced8190>
>>> for character in session.query(Character).all():
...     print(character.signature_phrase)
...
このきたならしい阿呆がァーッ!!
このクサレ脳ミソがァーーッ
このきたならしい阿呆がァーッ!!
このきたならしい阿呆がァーッ!!
このド低能がァーーッ
このクサレ脳ミソがァーーッ
このきたならしい阿呆がァーッ!!
このきたならしい阿呆がァーッ!!
このクサレ脳ミソがァーーッ
このド低能がァーーッ

もっと手の込んだデータを作りたいなら、BaseFuzzyAttributeクラスを継承して、 Custom fuzzy fields を作ることもできます。

最後に

factory_boyは実際に私が携わっている仕事で採用しているのですが、fixtureファイルを使っていた時にやっていた、他のテストのfixtureファイルをコピペで持ってきて、ちょっと編集して…のような作業が一切なくなって、かなり作業が楽になりました。

SQLAlchemy以外にも Django , Mogo , MongoEngine に対応しているので、いろいろなプロジェクトで使えるのではないでしょうか? Pythonでテストコードを書く際は、ぜひ使ってみてください。

明日は、 Attsun1031 さんです!