【boto3】DynamoDBの自動テストで使えるかもしれないスニペット

LocalStackのようなエミュレータを活用したテストコードでよく使うやつ。

• Python 3.8.3
• boto3 1.16.14
• pytest 6.1.2

ここで記載するサンプルメソッドは下のようなクラスに所属している前提になります。

import boto3
import mypy_boto3_dynamodb.service_resource as dynamodb_resources

class DynamoDefiner:
    dynamo: dynamodb_resources.DynamoDBServiceResource

    def __init__(self, endpoint_url: str = None):
        dynamo: dynamodb_resources.DynamoDBServiceResource = boto3.resource(
            "dynamodb", endpoint_url=endpoint_url
        )
        self.dynamo = dynamo

テーブルがない場合は作成

def create_my_table(self, table_name: str) -> None:
    # すでにテーブルがあるかチェック (二重にテーブルを作成しようとするとエラーになるため)
    tables = list(self.dynamo.tables.all().__iter__())
    if list(filter(lambda t: t.name == table_name, tables)):
        # すでにある場合は終了
        return

    self.dynamo.create_table(
        TableName=table_name,
        KeySchema=[
            {"AttributeName": "partition_key", "KeyType": "HASH"},
            {"AttributeName": "sort_key", "KeyType": "RANGE"},
        ],
        AttributeDefinitions=[
            {"AttributeName": "partition_key", "AttributeType": "S"},
            {"AttributeName": "sort_key", "AttributeType": "S"},
        ],
        BillingMode="PAY_PER_REQUEST",
        # GSI貼る場合
        GlobalSecondaryIndexes=[
            {
                "IndexName": "sort_key-index",
                "KeySchema": [
                    {"AttributeName": "sort_key", "KeyType": "HASH"},
                ],
                "Projection": {
                    "ProjectionType": "ALL",
                },
            },
        ],
    )

テーブルのItemを全件削除

def clear_my_table(self, table_name: str) -> None:
    my_table = self.dynamo.Table(table_name)
    # scanでItems取得
    scanning_result = my_table.scan()
    items = scanning_result["Items"]

    # 0件なら終了
    if not items:
        return

    # scanで取得したItemsを全件削除
    # TODO batch_writeに変えたら高速化できそう
    for item in items:
        shop_table.delete_item(
            Key={"shop_id": item["shop_id"], "sort_key": item["sort_key"]}
        )

    # 再帰
    self.clear_shop_table(table_name)

1回のscanで全件取得できるとは限らないため、取得しきれなかった場合は再帰呼び出しで繰り返すようにしています。

使い方

テストフレームワークはpytest、DynamoエミュレータはLocalStackを使用しているものとします。
setupでテスト毎に

• テーブルがなければ作成
• テーブルのクリア

をします。

import boto3
import mypy_boto3_dynamodb.service_resource as dynamodb_resources

class TestShopDynamoRepository(object):
    def setup_method(self) -> None:
        dynamoDefiner = DynamoDefiner(endpoint_url="http://localhost:4566")
        dynamoDefiner.create_shop_table("my-table")
        dynamoDefiner.clear_shop_table("my-table")

こうすると前のテストの状態に依存しないように各テストを実施できます。

参考

Python で DynamoDBを使う - Qiita