AWSのサービスをモックするライブラリmotoを拡張してDynamoDBのTransactWriteItemsを実装する

はじめに

サーバーレス開発部@大阪の岩田です。 現在開発中のサーバーレスアプリでDynamoDBのTransactWriteItemsを利用したいシーンが出てきたのですが、利用にあたって少し課題が出てきました。 DynamoDBのTransactWriteItemsは比較的新しい機能なので、ユニットテストに使用しているライブラリmotoが対応していないのです。

TransactWriteItems導入前はmotoでDynamoDBの各種操作をモックすることでうまくユニットテストが回っていたのですが、TransactWriteItemsを導入後にユニットテストを実行すると

...略
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <moto.dynamodb2.responses.DynamoHandler object at 0x10595e630>

    @amzn_request_id
    def call_action(self):
        self.body = json.loads(self.body or '{}')
        endpoint = self.get_endpoint_name(self.headers)
        if endpoint:
            endpoint = camelcase_to_underscores(endpoint)
>           response = getattr(self, endpoint)()
E           AttributeError: 'DynamoHandler' object has no attribute 'transact_write_items'

../../../../.venv/lib/python3.6/site-packages/moto/dynamodb2/responses.py:64: AttributeError

というエラーが出てユニットテストがこけるようになってしまいました。 transact_write_itemsなんて知らねーよ!!とのことです。

対策として自力でmotoを拡張しDynamoDBのTransactWriteItemsをモックできるようにしたので、手順をご紹介します。

motoとは?

freezegunの開発者として有名なSteve Pulec氏によって開発された、AWSの各種サービスをモックするためのPythonのライブラリです。 GitHubのスター数も2,000オーバーと数多くのユーザーに利用されています。

Python前提の話にはなりますが、DynamoDB LocalやLocalStackのようにいちいちモック用のサーバーを起動しなくてもモックできるのが便利な点で、DynamoDBを例に挙げると@mock_dynamodb()のデコレータをテスト対象に付与するだけでboto3の各種メソッドがmotoのモック処理に差し替わります。 CI/CDの中でDynamoDB LocalやLocalStackのDockerコンテナを立ち上げてユニットテストを回すと「何か良く分からないけどタイムアウトでテストが落ちた。やり直したら通った。。」なんてことも起こったりしますがmotoだとこういった問題は発生しません。

また、スタンドアロンのサーバーモードで動作させることも可能で、この場合はLocalStackのような感覚でPython以外のプログラミング言語やAWS CLIからもアクセスすることが可能です。

実装してみる

ここから実際にTransactWriteItemsを自力で実装していきます。

大枠の方針として、motoにMockのパッチを当てることでTransactWriteItemsを呼び出せるようにします。 本来はmoto本体を修正してプルリクを上げるのがベストなのですが、motoは開発が停滞気味でプルリクがマージされるまでに結構な時間がかかってしまいます。 私も過去に何度かバグ修正や機能追加を行ってプルリクを上げたことがあるのですが、大体マージされるまでに1ヶ月以上はかかっていました。 今もマージ待ちのプルリクが1つ・・・

さすがに1ヶ月ものんびり待ってられないので、自分の力だけで完結できる方式で進めていきます。 また、TransactWriteItemsの全機能を完璧に実装するのではなく、今回の案件で自分が利用する機能に限定して実装していきます。

motoのtransact_write_itemsにパッチする

まずテスト対象にデコレータを追加します。

@mock.patch('moto.dynamodb2.responses.DynamoHandler.transact_write_items',
                new=mock_transact_write_items, create=True)
    def test_for_blog(self):
        ...略

これでDynamoDBモックオブジェクトのtransact_write_itemsメソッドが呼び出された時に、mock_transact_write_itemsが呼び出されるようになります。

mock_transact_write_itemsの実装

次にmock_transact_write_itemsを実装します。 最終形は下記の通りです。

from moto.dynamodb2.responses import dynamo_json_dump


def mock_transact_write_items(self):

    def put_item(item):
        name = item['TableName']
        record = item['Item']
        # 使ってないパラメータは諸々省略
        expected = None
        overwrite = None
        return self.dynamodb_backend.put_item(name, record, expected, overwrite)

    def delete_item(item):
        name = item['TableName']
        keys = item['Key']
        # 使ってないパラメータは諸々省略
        return self.dynamodb_backend.delete_item(name, keys)

    def update_item(item):
        name = item['TableName']
        key = item['Key']
        update_expression = item.get('UpdateExpression')
        attribute_updates = item.get('AttributeUpdates')
        expression_attribute_names = item.get(
            'ExpressionAttributeNames', {})
        expression_attribute_values = item.get(
            'ExpressionAttributeValues', {})

        # 使ってないパラメータは諸々省略
        expected = None

        return self.dynamodb_backend.update_item(
            name, key, update_expression, attribute_updates, expression_attribute_names,
            expression_attribute_values, expected
        )

    transact_items = self.body['TransactItems']

    for transact_item in transact_items:
        if 'Put' in transact_item:
            put_item(transact_item['Put'])
        elif 'Update' in transact_item:
            update_item(transact_item['Update'])
        elif 'Delete' in transact_item:
            delete_item(transact_item['Delete'])

    return dynamo_json_dump({})

詳細を少しづつ見ていきます。 まず今回の案件では下記のようにtransact_write_itemsを呼び出しています。

dynamo.transact_write_items(
    TransactItems=[
        {
            'Update': {
                'TableName': 'test-table',
                'Key': {
                    'id': {'S': 'id1'}
                },
                'UpdateExpression': 'SET attr1 = :attr1',
                'ExpressionAttributeValues': {
                    ':attr1': {'S': '1234567'},
                }
            },
        },
        {
            'Put': {
                'TableName': 'test-table',
                'Item': {
                    'id': {'S': '1234567'},
                    'id_ref': {'S': 'id1'}
                },
                'ConditionExpression': 'attribute_not_exists(id)'
            },
        },
        {
            'Delete': {
                'TableName': 'test-table',
                'Key': {
                    'id': {'S': '123456'}
                },
            },
        }
    ])

re:Invent2018のDAT374で紹介されていたTransactWriteItemsのユースケース「User profile management 」と同じようなパターンです。

[レポート] DAT374 :[新機能!] Amazon DynamoDBのTransactionsを用いたモダンなアプリケーションの構築 #reinvent

この呼び出し方(呼び出され方)に対応するように、一連のトランザクション処理から1つづつテーブル操作を取り出してPUT,UPDATE,DELETEに振り分けるようにモックを実装します。

    transact_items = self.body['TransactItems']

    for transact_item in transact_items:
        if 'Put' in transact_item:
            put_item(transact_item['Put'])
        elif 'Update' in transact_item:
            update_item(transact_item['Update'])
        elif 'Delete' in transact_item:
            delete_item(transact_item['Delete'])

次にPUT,UPDATE,DELETEそれぞれの詳細を実装します。 boto3のput_item,delete_item,update_itemそれぞれのモック処理がmotoに実装済みなのでmoto.dynamodb2.responses.DynamoHandlerクラスのインスタンス変数self.dynamodb_backendからmoto.dynamodb2.models.DynamoDBBackendクラスの同名メソッドを呼び出すことでPUT,UPDATE,DELETEを実現します。 本来はパラメータのバリデーションやデフォルト値の設定、パラメータに応じた処理の分岐なども必要になるのですが、今回は自分が利用する機能に限定してモックを実装しています。

    def put_item(item):
        name = item['TableName']
        record = item['Item']
        # 使ってないパラメータは諸々省略
        expected = None
        overwrite = None
        return self.dynamodb_backend.put_item(name, record, expected, overwrite)

    def delete_item(item):
        name = item['TableName']
        keys = item['Key']
        # 使ってないパラメータは諸々省略
        return self.dynamodb_backend.delete_item(name, keys)

    def update_item(item):
        name = item['TableName']
        key = item['Key']
        update_expression = item.get('UpdateExpression')
        attribute_updates = item.get('AttributeUpdates')
        expression_attribute_names = item.get(
            'ExpressionAttributeNames', {})
        expression_attribute_values = item.get(
            'ExpressionAttributeValues', {})

        # 使ってないパラメータは諸々省略
        expected = None

        return self.dynamodb_backend.update_item(
            name, key, update_expression, attribute_updates, expression_attribute_names,
            expression_attribute_values, expected
        )

試してみる

実装できたのでテストを実行します。

============================= test session starts ==============================
platform darwin -- Python 3.6.1, pytest-4.1.0, py-1.7.0, pluggy-0.8.1
rootdir: /Users/iwata.tomoya/xxxxxxx, inifile:collected 1 item

test_for_blog.py                                                        [100%]

=========================== 1 passed in 3.15 seconds ===========================.
Process finished with exit code 0

今度は正常に終了しました!! 詳細は割愛していますが、transact_write_itemsを呼び出した後に再度get_itemしてassertするテストが問題なく通るようになりました。

まとめ

motoを拡張してDynamoDBのTransactWriteItemsを実装する手順について見て来ました。 同様の手順で他のAWSサービスやDynamoDBの他のAPIについてもモックを実装できるので、積極的にAWSの新機能を取り入れて導入していきたいですね!!

DynamoDBのTransactWriteItemsとTransactGetItemsに関しては落ち着いたらちゃんと実装してmoto本体にプルリクを上げたいと思います。