【moto】motoでtransact_write_itemsをモックする。【Python】

motoでtransact_write_itemsをモックする。

moto と transact_write_items で検索すると、motoを拡張してモックする方法が上位に表示される。詳しく調べていないが、moto 2.2.4 では拡張せずともモックできるので、サンプルを提示したい。

この記事ではテストコードは一部だけ記載するので、完全なものは https://github.com/tomohikoseven/moto_transact_write_items を参照ください。

実装コード

実装したコードは、ClickCountというテーブルにItemをトランザクションで追加・更新（アトミックカウンタ）するものです。DynamoDBをご存知の方には、簡単なものかと思います。

簡単に説明すると、以下の実装コードはハッシュキー：month は文字列 YYYYMM を、レンジキー（ソートキー）：date は文字列 YYYYMMDD を期待して、1トランザクションで2つのItemに対し、clickCountをカウントUPするものになります。

トランザクションの1つ目は日付の末尾が00となっており、月間のclickCountをカウントUPします。今回の記事では気にすることではないです。

(index.py)
import boto3


client = boto3.client('dynamodb',  region_name='ap-northeast-1')
def lambda_handler( month:str, date:str ):
  try:
    response = client.transact_write_items(
      TransactItems=[
        {
          'Update' : {
            'Key' : {
              'hKey': { 'S': month },
              'rKey': { 'S': month + '00' }
            },
            'TableName':'ClickCount',
            'UpdateExpression': 'ADD clickCount :increment',
            'ExpressionAttributeValues' : {
              ':increment': { 'N': '1' }
            }
          }
        },
        {
          'Update' : {
            'Key' : {
              'hKey': { 'S': month },
              'rKey': { 'S': date }
            },
            'TableName':'ClickCount',
            'UpdateExpression': 'ADD clickCount :increment',
            'ExpressionAttributeValues' : {
              ':increment': { 'N': '1' }
            }
          }
        }
      ]
    )
  except Exception as e :
    print("===")
    print(e)
    raise e

  return 200

テストコード

pytestでテストコードを実行します。

motoでモックするテストコードを記述したことがある方は難しいことはないです。mock_dynamodb2をimportして、アノテーション@mock_dynamodb2をテストケースに付けてやればよいです。

(test_index.py)

# 本当のテストは1ケースだけではないので、ここでは余分なライブラリがインポートされています。
from importlib import import_module
import boto3
from moto import mock_dynamodb2
from unittest import mock
from botocore.exceptions import ClientError
import pytest


@mock_dynamodb2
def test_success(set_dynamodb_Count):
  # mock dynamodb
  table = set_dynamodb_Count()

  month = '202101'
  date = '20210101'
  from src.index import lambda_handler
  res = lambda_handler( month, date )

  item = table.get_item(Key={'hKey':month, 'rKey': date})
  assert res == 200
  assert item['Item']['clickCount'] == 1

conftest.pyに set_dynamodb_Count を記載しています。テストにはこちらも必要になります。test_index.pyと同一フォルダに置いてください。

(conftest.py)
import pytest
import boto3

@pytest.fixture
def set_dynamodb_Count():
  def definition():
    dynamodb = boto3.resource('dynamodb', region_name='ap-northeast-1')
    dynamodb.create_table(
      TableName='ClickCount',
      KeySchema=[
          {
              'AttributeName': 'hKey',
              'KeyType': 'HASH'
          },
          {
              'AttributeName': 'rKey',
              'KeyType': 'RANGE'
          }
      ],
      AttributeDefinitions=[
          {
              'AttributeName': 'hKey',
              'AttributeType': 'S'
          },
          {
              'AttributeName': 'rKey',
              'AttributeType': 'S'
          }
      ],
    )

    return dynamodb.Table('ClickCount')
  return definition