Boto3のMockのMotoに未実装のAPIを追加してみた

きっかけ

AWS Lambdaのユニットテストって皆さんどのように実装しているのでしょうか？
私は前回の記事でも少し登場していたMotoを使用しています。テストケースにデコレータを付けるだけでBoto3をモック化してくれるのでとても便利なのですが、先日ついに恐れていたことが起きてしまいました。

NotImplementedError: The create_group action has not been implemented

cognito-idpのcreate_group()を使用した関数をテストしたら「残念だったな、そのAPIはまだ未実装だぜ！」と言われてしまったのです。（口調は私の想像）

こんな時には（通常）のMockでどうにかするのが良いと思うのですが、私のPython力不足でMockの使い方がよく分からず、足掻きに足掻いた結果、Boto3のMockのMotoの関数をMock化することでどうにか解決できたので、その方法をまとめておきます。

問題のテストケース

テストケースと言うよりもエラーが出た処理なのですが、単にCognitoユーザープールを作成して、それにグループを作成するだけです。create_user_pool()は問題なく動きますが、create_group()が未実装のためエラーとなります。

import boto3
from moto import mock_cognitoidp

@mock_cognitoidp
def cognito_mock_test():
    cognito = boto3.client('cognito-idp')
    response = cognito.create_user_pool(
        PoolName='TestUserPool'
    )
    user_pool_id = response['UserPool']['Id']
    response = cognito.create_group(  # ここでNotImplementedError
        GroupName='TestGroup',
        UserPoolId=user_pool_id
    )

create_group()のMock化

CognitoIdpResponseというクラスに各関数のMockを実装しているので、そこにcreate_group()をpatchしてあげます。これでcreate_group()を呼び出した時には、代わりにmock_create_group()が呼び出されるようになります。元のCognitoIdpResponseにはcreate_group()が存在しないので、create=Trueを指定する必要があります。

from unittest import mock

@mock_cognitoidp
@mock.patch('moto.cognitoidp.responses.CognitoIdpResponse.create_group',
            new=mock_create_group, create=True)
def cognito_mock_test():
    # 上と同じ

そして、mock_create_group()の実装が以下です。
self.bodyに引数が入っています。戻り値は3つ指定する必要があり、順にステータスコード、追加のヘッダー、コンテンツです。

import json

def mock_create_group(self):
    group_name = json.loads(self.body)['GroupName']
    return '200', {}, json.dumps({
        'Group': {
            'GroupName': group_name
        }
    })

本物のcreate_group()はGroupName以外にもコンテンツがありますが、ひとまずこれだけにしておきます。

{'Group': {'GroupName': 'TestGroup'},
 'ResponseMetadata': {'HTTPHeaders': {'server': 'amazon.com'},
                      'HTTPStatusCode': 200,
                      'RetryAttempts': 0}}

これで元の関数を実行すると、見事にcreate_group()がレスポンスを返してくれるようになりました！

オチ

こんな感じでMotoに不足していたAPIをMock化してみたのですが、Motoが1.3.8にアップデートされ、create_group()が実装されちゃっていました！ｗ　本家のレスポンスはコチラ。

{'Group': {'CreationDate': datetime.datetime(2019, 4, 30, 19, 41, 39, tzinfo=tzlocal()),
           'Description': '',
           'GroupName': 'TestGroup',
           'LastModifiedDate': datetime.datetime(2019, 4, 30, 19, 41, 39, tzinfo=tzlocal()),
           'UserPoolId': 'ap-northeast-1_717df8345675458c98e2b9190a1fc580'},
 'ResponseMetadata': {'HTTPHeaders': {'server': 'amazon.com'},
                      'HTTPStatusCode': 200,
                      'RetryAttempts': 0}}

Mock実装は無駄になってしまいましたが、色々勉強になったし、今後も未実装のAPIが出てきた時には使えそうなので、これはこれで良かったのかな。

ちなみにMotoが1.3.8になって、依存性で今までaws-xray-sdk==0.95が入っていたのが、ついにaws-xray-sdk==2.4.2（最新）になりました。早速プロジェクトで使っているものもバージョンアップしなければ！