Python: freezegun で時刻のテストを楽に書く

時刻周りの処理はバグが混入しやすい上にテストが書きづらくて面倒くさい。今回は、そんな面倒な時刻のテストを楽に書けるようになる freezegun というパッケージを使ってみる。この freezegun というパッケージを使うと Python の標準ライブラリの datetime から得られる現在時刻を指定したものに差し替えることができる。

使った環境は次の通り。

$ sw_vers
ProductName:    Mac OS X
ProductVersion: 10.11.6
BuildVersion:   15G1108
$ python --version
Python 3.5.2

時刻をテストするときの面倒くささ

時刻周りの処理をテストをするときは、当然ながら色々な時刻を使ってテストがしたい。とはいえ、そのためだけにシステムの時刻を変更しながらテストを走らせるわけにもいかないだろう。そこで、大抵の場合は時刻を取得するところをモックに置き換えてテストすることになるはず。ただ、Python の datetime はモックアウトが意外とやりにくい。

これは何故かというと datetime の提供する関数がちょいちょい C 拡張モジュールとして書かれているため。 C 拡張モジュールで書かれていると標準ライブラリの unittest.mock で置き換えることができない。

例えば、次のように now() 関数をピンポイントでモックに置き換えようとすると例外になってしまう。

>>> from unittest import mock
>>> with mock.patch('datetime.datetime.now') as now:
...     now.return_value = datetime(2015, 10, 21)
...     from datetime import datetime
...     datetime.now()
...
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/amedama/.pyenv/versions/3.5.2/lib/python3.5/unittest/mock.py", line 1312, in __enter__
    setattr(self.target, self.attribute, new_attr)
TypeError: can't set attributes of built-in/extension type 'datetime.datetime'

そのため datetime オブジェクトを丸ごとモックに置き換える必要がある。

>>> from unittest import mock
>>> import datetime
>>> fake_dt = datetime.datetime(2015, 10, 21)
>>> with mock.patch('datetime.datetime') as dt:
...     dt.now.return_value = fake_dt
...     from datetime import datetime
...     datetime.now()
...
datetime.datetime(2015, 10, 21, 0, 0)

ただ、これだと datetime オブジェクトの別のメソッドを使いたいときなんかに困る。その都度、書き換えたり書き戻したりが必要になって、うーん…という感じ。

freezegun を使う

そこで freezegun を使うと、そういった手間がかからなくなる。

パッケージのインストールは pip コマンドから。

$ pip install freezegun

特定の時刻を返すようにする

freezegun では freezegun.freeze_time() で datetime モジュールの関数が特定の時刻を返すように置き換えることができる。

例えば以下はその API をデコレータとして使うパターン。デコレータで修飾された関数の中でのみ時刻が置き換わる。サンプルコードでは datetime.now() の内容を見ている。

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import freezegun
from datetime import datetime


@freezegun.freeze_time('2015-10-21')
def main():
    print(datetime.now())


if __name__ == '__main__':
    main()

上記を実行してみよう。

$ python freeze.py
2015-10-21 00:00:00

見事に取得できる時刻が置き換わっている。ちなみに、この日付はバック・トゥ・ザ・フューチャーのアレ。

上記では日付までしか指定しなかったけど、もちろん時刻まで入れられる。

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import freezegun
from datetime import datetime


def main():
    with freezegun.freeze_time('2015-10-21 12:34:56'):
        print(datetime.now())


if __name__ == '__main__':
    main()

上記の実行結果は次の通り。

$ python freeze.py
2015-10-21 12:34:56

ちなみに、引数には置き換えたい時刻の入った datetime オブジェクトを渡しても良い。

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import freezegun
from datetime import datetime


dt = datetime(2015, 10, 21, 12, 34, 56)


@freezegun.freeze_time(dt)
def main():
    print(datetime.now())


if __name__ == '__main__':
    main()

結果については先ほどと同じ。

特定の時刻を返すようにする (コンテキストマネージャ)

先ほどはデコレータを使って時刻を指定した。同じ API はコンテキストマネージャとしても使うことができる。

以下はコンテキストマネージャとして使ったパターン。この場合は、コンテキストマネージャのブロック内でだけ時刻が置き換わる。

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import freezegun
from datetime import datetime


def main():
    with freezegun.freeze_time('2015-10-21'):
        print(datetime.now())


if __name__ == '__main__':
    main()

実行結果についてはデコレータを使った場合と変わらないので省略する。

特定の時刻を返すようにする (素)

デコレータとしてもコンテキストマネージャとしても使わない場合は、こんな感じ。とはいえ、まあこれはやらないかな。

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import freezegun
from datetime import datetime


def main():
    freezer = freezegun.freeze_time('2015-10-21')
    freezer.start()
    try:
        print(datetime.now())
    finally:
        freezer.stop()


if __name__ == '__main__':
    main()

実行結果は同じ。

特定の時刻から時間をずらしていく

テストをするときは一旦基準となる時刻に規正してから、特定の処理をした後にずらしたいというニーズもあると思う。 freezegun を使えば、それも簡単にできる。

次のサンプルコードでは tick() を使って時刻を timedelta オブジェクトの分だけずらしたり move_to() を使って時刻を切り替えている。

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import freezegun
from datetime import datetime
from datetime import timedelta


def main():
    with freezegun.freeze_time('2015-10-21 00:00:00') as freeze_datetime:
        print(datetime.now())

        # 時間を 1 秒進める
        freeze_datetime.tick()
        print(datetime.now())

        # 時間を 1 分進める
        freeze_datetime.tick(delta=timedelta(minutes=1))
        print(datetime.now())

        # 特定の時間に移す
        freeze_datetime.move_to('2000-01-01 00:00:00')
        print(datetime.now())


if __name__ == '__main__':
    main()

上記を実行すると、次のようになる。

$ python freeze.py
2015-10-21 00:00:00
2015-10-21 00:00:01
2015-10-21 00:01:01
2000-01-01 00:00:00

ユニットテストで使う

ここまでで一通りの使い方は紹介できたので、あとはユニットテストに組み込んで使うだけ。とはいえ、何も難しいことはない。

例えば標準ライブラリの unittest を使ったテストで使うなら、こんな感じ。

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import unittest

import freezegun
from datetime import datetime


class Test(unittest.TestCase):

    @freezegun.freeze_time('2015-10-21')
    def test(self):
        self.assertEqual(datetime.now(), datetime(2015, 10, 21))


if __name__ == '__main__':
    unittest.main()

実行してみよう。

$ python test_datetime.py -v
test (__main__.Test) ... ok

----------------------------------------------------------------------
Ran 1 test in 0.006s

OK

ばっちりだね。