実戦TDD! / PythonのMagicMockを使ってDynamoDBをモックしたテストコードを実装する

今は昔、チームメンバーがPython3で書いたコードありけり。コード、AWS Lambdaにデプロイされて動作しつつ、ときたまエラーを吐きてバグりけり。

 

……しかし、このPythonプロジェクト、テストコードがない(古文調は諦めた)。

 

そこで、ただバグを直しても芸がないと思い、ついでにPythonのテストコードの書き方を学びつつ、テスト駆動開発(TDD)を実戦投入してみることにした。

以下、Python初心者によるTDDの記録である。

 

 

[DynamoDBへのアクセス時エラーをテストする]

問題のエラーはAmazon DynamoDBにアクセスする際に発生しており、エラー処理が適切に行われていなかった。

 

このエラー処理のテストを書くには、3つの選択肢がある。

  1. 実物のDynamoDBを利用する
  2. DynamoDBLocalやLocalstackなどDynamoDB APIと互換性のあるミドルウェアを利用する
  3. テストライブラリを利用してDynamoDBへのアクセスをモックする

 

今回は選択肢3を採用した。理由は以下の通り。

  • Pythonでテストコードを書いたことがなかったので勉強したかった
  • モックならネットワークアクセスが発生しないのでテスト実行が比較的高速
  • ミドルウェアや実行環境の差異による影響をもっとも受けにくい

 

選択肢1については、以下の理由で断念した。

  • DynamoDBへのアクセス権限を各開発者ローカル環境やCI環境毎にセットアップするのが面倒
  • 微々たる額ながらDynamoDBの利用料金がかかる
  • インターネット経由でのDynamoDBへのアクセスによりテスト実行にかかる時間が長くなる

 

選択肢2については、以下の理由で断念した。

 

[MagicMockでDynamoDBへのアクセスをモックする]

Pythonでテストコードを書いたことがなかったので、適当に検索して公式ドキュメントにたどり着いた。このモックオブジェクトライブラリーが利用できるのはPython 3.3以降らしい。

docs.python.org

 

MagicMockクラスを使えば、どの関数でもモックすることができる。ということは、これを使えばDynamoDBに限らず、boto3が対応している他のAWSサービスへのアクセスも同じようなやり方でテストを記述することができる。

 

ローカルでの動作確認は今回Python 3.7.1を利用した。

$ python --version
Python 3.7.1

 

バグを修正する前のテスト対象のファイルは、擬似的な表すと次のようになっていた。

[myfunc.py]

import boto3

def lambda_handler(event, context):
    table = boto3.resource('dynamodb').Table('my_awesome_table')
    update(table, event['id'], event['count_up'])

def update(table, id, count_up):
# call DynamoDB API table.update_item( Key={ 'id': id }, UpdateExpression='SET #count = #count + :count_up', ExpressionAttributeNames={ '#count': 'count' }, ExpressionAttributeValues={ ':count_up': count_up } ) return True # Success

本当はハンドラ関数()の外でtableを初期化する方がパフォーマンスが良いのだが、このラムダ関数にはパフォーマンスはあまり求められていないので無視した。

 

これをモックしたテストコードを書きはじめる。とりあえずはDynamoDBのUpdateItemが呼ばれていることをアサートする。

[test_myfunc.py]

import myfunc
import boto3
import unittest
from unittest import TestCase
from unittest.mock import MagicMock

table = boto3.resource('dynamodb').Table('my_awesome_table')

class TestUpdate(TestCase):
    def setUp(self):
        # prepare the test
        table.update_item = MagicMock()
    
    def test_update_normal(self):
        # call the function calling DynamoDB API
        myfunc.update(table, 'dummy_id', 1)
        
        # assert DynamoDB API is called
        table.update_item.assert_called_once()
    
    def tearDown(self):
        # clean up

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

 

[モックした関数呼び出しをテストする]

ラムダ関数本体のコードとテストコードは、以下のように同じディレクトリに置かれている。

$ ls
myfunc.py	test_myfunc.py

 

このとき、Pythonのunittestを実行するには、python -m unittest discoverというコマンドを実行すれば良い。しかし、普段Pythonを使い慣れていないので、この長ったらしいコマンドをこの先もずっと覚えていられる気がしない。 

 

そこで、以下のようにMakeファイルを書いて、make testコマンドでテストを実行できるようにしておいた。

[Makefile]

.DEFAULT_GOAL := help

help: ## show help text
    @echo "Description:"
    @echo "    Run tests."
    @echo "Commands:"
    @echo "    test    Execute unittest"

test:
    python -m unittest discover

 

コマンドを実行すると以下のようにテストが実行され、結果が表示される。

$ make test
python -m unittest discover
.
---------------------------
Ran 1 tests in 0.05s

OK

無事にテストが通った。

 

[MagicMockのアサートを利用する]

先ほど利用したassert_called_once()以外にも、MagicMockにはいくつかのアサート用関数が用意されている。

  • assert_called(*args, **kwargs):最低1回はモックが呼び出されたことをアサートする
  • assert_called_once_with(*args, **kwargs):指定の引数で1回だけモックが呼び出されたことをアサートする(2回以上呼ばれた場合はテストが失敗する)
  • assert_not_called():モックが呼び出されていないことをアサートする

 

また、MagicMockにはcall_countというプロパティがある。これを利用することで、次のようにモックの呼び出し回数をアサートすることができる。

# import文は省略

table = boto3.resource('dynamodb').Table('my_awesome_table')

class MyTest(TestCase):
    def test_mock_call_count(self):
        table.update_item = MagicMock()
        table.update_item()
        table.update_item()
        table.update_item()
    
        self.assertEqual(table.update_item.call_count, 3) # True

 

これらを活用して、次のような形でいくつかの引数のパターンのテストを拡充していく。

# 色々省略。
    def test_normal_argument(self):
        table.query = MagicMock()
        table.update_item = MagicMock()
        
        myfunc.update(table, 'new_test_id', 100)
        
        table.query.assert_not_called()
        table.update_item.assert_called_once_with(
            Key={
                'id': 'new_test_id'
            },
            UpdateExpression='SET #count = #count + :count_up',
            ExpressionAttributeNames={
                '#count': 'count'
            },
            ExpressionAttributeValues={
                ':count_up': 100
            }
        )

 

[引数リストのアンパックを利用してリファクタ]

順調にテストケースを追加していったところ、下記のように、似たような引数での関数呼び出しに対するアサートがいくつも並んでしまった。

# 色々省略。
    def test_2(self):
        # 色々省略。
        table.update_item.assert_called_once_with(
            Key={
                'id': 'test_2'
            },
            UpdateExpression='SET #count = #count + :count_up',
            ExpressionAttributeNames={
                '#count': 'count'
            },
            ExpressionAttributeValues={
                ':count_up': 2
            }
        )

    def test_3(self):
        # 色々省略。
        table.update_item.assert_called_once_with(
            Key={
                'id': 'test_3'
            },
            UpdateExpression='SET #count = #count + :count_up',
            ExpressionAttributeNames={
                '#count': 'count'
            },
            ExpressionAttributeValues={
                ':count_up': 3
            }
    
    def test_4(self):
        # 色々省略。
        table.update_item.assert_called_once_with(
            Key={
                'id': 'test_4'
            },
            UpdateExpression='SET #count = #count + :count_up',
            ExpressionAttributeNames={
                '#count': 'count'
            },
            ExpressionAttributeValues={
                ':count_up': 4
            }

 

これでは、例えばtest_5のような新しいテストケースを作る際に、コピペが発生してしまう。

 

何かいい方法はないかと考え、Google先生に尋ねたところ(検索キーワードは忘れてしまったが)、以下の公式ドキュメントにたどりついた*1

docs.python.org

公式ドキュメントで上げられている例に少し追加して使い方をメモしておく。

 

キーワード引数でない通常の引数をアンパックする場合には、*演算子を用いる。

>>> list(range(3, 6))    # 別々の引数での通常呼び出し
[3, 4, 5]
>>> args = [3, 6]
>>> list(range(*args))    # リストからアンパックされた引数での呼び出し
[3, 4, 5]

キーワード引数をアンパックする場合には、**演算子を用いる。

>>> def parrot(voltage, state='a stiff', action='voom'):
...     print("-- This parrot wouldn't", action, end=' ')
...     print("if you put", voltage, "volts through it.", end=' ')
...     print("E's", state, "!")
...
>>> parrot(voltage="three million")    # 別々のキーワード引数での通常呼び出し
-- This parrot wouldn't voom if you put three million volts through it. E's a stiff !
>>> d = {"voltage": "four million", "state": "bleedin' demised", "action": "VOOM"}
>>> parrot(**d)    # 辞書からアンパックされたキーワード引数での呼び出し
-- This parrot wouldn't VOOM if you put four million volts through it. E's bleedin' demised !

 

このキーワード引数のアンパックを利用して、先ほどのコードを次のようにリファクタした。

# 色々省略。
def get_args(count_up):
    return {
        'Key': {
            'id': 'test_' + str(count_up)
        },
        'UpdateExpression': 'SET #count = #count + :count_up',
        'ExpressionAttributeNames': {
            '#count': 'count'
        },
        'ExpressionAttributeValues': {
            ':count_up': count_up
        }
    }

class MyTest(TestCase):
    def test_2(self):
        # 色々省略。
        table.update_item.assert_called_once_with(**get_args(2))
    
    def test_3(self):
        # 色々省略。
        table.update_item.assert_called_once_with(**get_args(3))
        
    def test_4(self):
        # 色々省略。
        table.update_item.assert_called_once_with(**get_args(4))

 

この後、本題であるエラー発生時のテストコードを書き、それに合わせてLambda関数本体のコードを修正した。その後、同様のエラーの発生は確認されていない。

 

[まとめ]

  • Python3.3以降では標準のモックオブジェクトライブラリ unittest.mock が使える
  • とりあえず MagicMock を使っておくと良い
  • 同一あるいは類似の引数での関数呼び出しを複数回アサートする際は * 演算子** 演算子で引数リストをアンパックすると便利

 

力尽きたので長くなってしまったので、この記事はここまでにして、TDDの本題であるエラー処理とそのテストについては、後日別記事を書く予定。

*1:日本語版のページを見たときに今回参照した「4.7.4. 引数リストのアンパック」が未翻訳だった。これはプルリクエストチャンスかと勇んだものの、CONTRIBUTING.mdでプルリクエストは受けつけていないと書かれているのを見つけてちょっとがっかり。GitHubに草は生えないけど、余裕ができたら翻訳にも貢献したい。