pytestとmotoを使ってPynamoDBの単体テストを実行してみる

はじめに

データアナリティクス事業本部のkobayashiです。

PythonでDynamoDBを操作する際にboto3を使っても良いのですが、Object-Document Mapper（ODM）として大変使いやすいPynamoDBがあります。今回はPynamoDBで記述したコードの単体テストをモックフレームワークのｍotoを使って行ってみたいと思います。

• Welcome to PynamoDB’s documentation! — PynamoDB 5.5.1 documentation

• getmoto/moto: A library that allows you to easily mock out tests based on AWS infrastructure.

PynamoDBとは

PynamoDBはDynamoDB用のODMになり、DynamoDBの 様々な処理を簡単にに記述できるインターフェースを提供しています。 PynamoDBは非常に使いやすく、直感的な記述を行うことができます。詳しくは当ブログエントリの以下の記事をご確認ください。

一方motoはAWSサービスをモック化するツールで、これを用いることでAWSの各種サービスをローカル環境で再現することができます。これは、テストや開発の際に非常に便利です。Motoを使うことで、APIの呼び出しを再現しながらも、実際のAWS環境に影響を与えずに、AWSの各サービスに対するコードのテストを行うことができます。こちらも詳しくは当ブログエントリの以下の記事をご確認ください。

motoでPynamoDBの単体テストを行ってみる

それでは実際にPynamoDBでDynamoDBのテーブルを使うコードを記述し、それに対するテストコードも作成してみます。

環境

• Python: 3.11.4
• pynamodb: 5.5.1
• moto: 4.2.12
• pytest: 7.4.3

はじめにPynamoDBでユーザーデータを保存するユーザーテーブルをモデル化します。

main.py

from pynamodb.models import Model
from pynamodb.attributes import UnicodeAttribute


class UserModel(Model):
    """
    A DynamoDB User
    """

    class Meta:
        table_name = 'dynamodb-user'
        region = 'us-west-1'

    email = UnicodeAttribute(hash_key=True)
    first_name = UnicodeAttribute()
    last_name = UnicodeAttribute()

    @classmethod
    def get_user(cls, email):
        return UserModel.get(email)

    @classmethod
    def create_user(cls, email, first_name, last_name):
        user = UserModel(email, first_name=first_name, last_name=last_name)
        user.save()

    @classmethod
    def get_users(cls, last_name):
        users = UserModel.scan(UserModel.last_name.startswith(last_name))
        return [user for user in users]

    @classmethod
    def update_user(cls, email, first_name, last_name):
        user = UserModel.get(email)
        user.update(
            actions=[
                UserModel.first_name.set(first_name),
                UserModel.last_name.set(last_name),
            ]
        )

        return UserModel.get(email)

    @classmethod
    def delete_user(cls, email):
        user = UserModel.get(email)
        user.delete()

PynamoDBのチュートリアルにあるコード を元にユーザーを取得・作成・更新・削除するクラスメソッドを追加してあります。

テーブルとしてはemailをパーティションキーとし、first_nameとlast_nameを属性にもったテーブルとなります。

ではこのユーザーモデルの単体テストを記述したいと思います。

以下が単体テストコードになります。

test_main.py

from moto import mock_dynamodb
import pytest

from main import UserModel

USER_DATA = [
    {"email": "test@example.com", "first_name": 'Samuel', "last_name": 'Adams'},
    {"email": "test_01@example.com", "first_name": 'Abigail', "last_name": 'Suzuki'},
    {"email": "test_02@example.com", "first_name": 'Diana', "last_name": 'Sato'},
    {"email": "test_03@example.com", "first_name": 'Lauren', "last_name": 'Tanaka'},
    {"email": "test_04@example.com", "first_name": 'Theresa', "last_name": 'Yamamoto'},
    {"email": "test_05@example.com", "first_name": 'Nicola', "last_name": 'Yamada'},

]


class TestUserModel:
    @pytest.fixture(scope="class")
    def pynamo_setup(self):
        """
        Usersテーブルのモック作成
        Returns:

        """
        with mock_dynamodb():
            UserModel.create_table(read_capacity_units=5, write_capacity_units=5, wait=True)

            # 初期データの登録
            for _u in USER_DATA:
                user = UserModel(**_u)
                user.save()

            yield

    @pytest.mark.parametrize(
        [
            "data_in", "expected_data"
        ],
        [
            pytest.param("test@example.com",
                         {"email": "test@example.com", "first_name": 'Samuel', "last_name": 'Adams'}, ),
            pytest.param("test_01@example.com",
                         {"email": "test_01@example.com", "first_name": 'Abigail', "last_name": 'Suzuki'}, ),
            pytest.param("test_02@example.com",
                         {"email": "test_02@example.com", "first_name": 'Diana', "last_name": 'Sato'}, ),
            pytest.param("test_03@example.com",
                         {"email": "test_03@example.com", "first_name": 'Lauren', "last_name": 'Tanaka'}, ),
            pytest.param("test_04@example.com",
                         {"email": "test_04@example.com", "first_name": 'Theresa', "last_name": 'Yamamoto'}, ),
            pytest.param("test_05@example.com",
                         {"email": "test_05@example.com", "first_name": 'Nicola', "last_name": 'Yamada'}, ),

        ],
    )
    def test_get_user(self, pynamo_setup, data_in, expected_data):
        # ユーザーを取得する
        user = UserModel.get_user(data_in)

        # 確認
        assert user.email == expected_data["email"]
        assert user.first_name == expected_data["first_name"]
        assert user.last_name == expected_data["last_name"]

    @pytest.mark.parametrize(
        [
            "data_in"
        ],
        [
            pytest.param("nouser@example.com"),

        ],
    )
    def test_get_user_error(self, pynamo_setup, data_in):
        # 存在しないユーザーを取得する
        with pytest.raises(Exception) as err:
            UserModel.get_user(data_in)

        # 確認
        assert err.typename == 'DoesNotExist'
        assert err.value.msg == "Item does not exist"

    @pytest.mark.parametrize(
        [
            "data_in",
        ],
        [
            pytest.param({"email": "test_11@example.com", "first_name": "Misaki", "last_name": "Suzuki", }, ),
            pytest.param({"email": "test_12@example.com", "first_name": "Norihito", "last_name": "Baba", }, ),
            pytest.param({"email": "test_13@example.com", "first_name": "Aya", "last_name": "Kinoshita", }, ),
            pytest.param({"email": "test_14@example.com", "first_name": "Hiroaki", "last_name": "Hirosue", }, ),
            pytest.param({"email": "test_15@example.com", "first_name": "Atsushi", "last_name": "Suzuki", }, ),
        ],
    )
    def test_create_user(self, pynamo_setup, data_in):
        # ユーザーを作成する
        user = UserModel.create_user(**data_in)
        # ユーザーを取得する
        ret_user = UserModel.get_user(data_in["email"])

        # 確認
        assert ret_user.email == data_in["email"]
        assert ret_user.first_name == data_in["first_name"]
        assert ret_user.last_name == data_in["last_name"]

    def test_count_user(self, pynamo_setup):
        cnt = UserModel.count()

        # 確認
        assert cnt == 11

    def test_scan_user(self, pynamo_setup):
        # last_nameがSuzukiから始まるユーザーを取得
        users = UserModel.get_users("Suzuki")
        users = [user for user in users]

        # 確認
        assert len(list(users)) == 3
        assert 'Abigail' in [user.first_name for user in users]
        assert 'Misaki' in [user.first_name for user in users]
        assert 'Atsushi' in [user.first_name for user in users]

    @pytest.mark.parametrize(
        [
            "data_in",
        ],
        [
            pytest.param({"email": "test@example.com", "first_name": 'User', "last_name": 'hoge'}, ),
            pytest.param({"email": "test_01@example.com", "first_name": 'User_01', "last_name": 'hoge_01'}, ),
            pytest.param({"email": "test_02@example.com", "first_name": 'User_02', "last_name": 'hoge_02'}, ),
            pytest.param({"email": "test_03@example.com", "first_name": 'User_03', "last_name": 'hoge_03'}, ),
            pytest.param({"email": "test_04@example.com", "first_name": 'User_04', "last_name": 'hoge_04'}, ),
            pytest.param({"email": "test_05@example.com", "first_name": 'User_05', "last_name": 'hoge_05'}, ),
            pytest.param({"email": "test_11@example.com", "first_name": "User_11", "last_name": 'hoge_11'}, ),
            pytest.param({"email": "test_12@example.com", "first_name": "User_12", "last_name": 'hoge_12'}, ),
            pytest.param({"email": "test_13@example.com", "first_name": "User_13", "last_name": 'hoge_13'}, ),
            pytest.param({"email": "test_14@example.com", "first_name": "User_14", "last_name": 'hoge_14'}, ),
            pytest.param({"email": "test_15@example.com", "first_name": "User_15", "last_name": 'hoge_15'}, ),
        ],
    )
    def test_update_user(self, pynamo_setup, data_in):
        # ユーザーを作成
        user = UserModel.update_user(**data_in)

        # 作成したユーザーを確認
        UserModel.get(data_in["email"])

        # 確認
        assert user.first_name == data_in["first_name"]

    @pytest.mark.parametrize(
        [
            "data_in"
        ],
        [
            pytest.param("test@example.com"),
            pytest.param("test_01@example.com"),
            pytest.param("test_11@example.com"),
        ],
    )
    def test_delete_user(self, pynamo_setup, data_in):
        # ユーザーを削除
        UserModel.delete_user(data_in)

        # 存在しないユーザーを取得
        with pytest.raises(Exception) as err:
            UserModel.get(data_in)

        assert err.typename == 'DoesNotExist'
        assert err.value.msg == "Item does not exist"

ポイントとしてはmotoのmoto_dyanmodbをwith句に使用したpytestのfixtureを作り、モックテーブルを作成後にyieldでテストを実行しています。

今回はテストクラスで1つのテーブルのCRUDを行いたいのでfixtureのscopeはclassにしています。

この状態で各テストでモックテーブルを使うには通常のfixtureの使い方と同じ様に引数に「Usersテーブルのモック作成」の関数名を追加するだけになります。

これでテストの記述は終わったのであとは実行してみます。

$ pytest main.py
=========================================================================== test session starts ============================================================================
platform darwin -- Python 3.11.4, pytest-7.4.3, pluggy-1.3.0
rootdir: /tmp/pynamodb-test
plugins: anyio-3.7.0
collected 0 items

このような形で実際のDynamoDBテーブルのデータを使用せずにテストを行うことができます。

まとめ

motoを使ってDynamoDBへの接続部分をモック化してPynamoDBの単体テストを行ってみました。fixtureを使って事前にテーブルを作成するといった一般的な方法で簡単にテストが行えることがわかりました。

最後まで読んで頂いてありがとうございました。