PynamoDBで良い感じにTestableなモデルを定義して、DynamoDB Localを使ってテストする方法

Python AWS

f:id:TatchNicolas:20201211101137p:plain

TL; DR;

  • PynamoDBを使ったテストでローカルで動かすDynamoDBを叩きたい場合に
  • テーブルのキー定義やキャパシティ設定の管理はterraform/CDKなどに任せつつ、アプリケーションコードを汚さずにテストを実行したい
  • そんなときは getattr でmetaclassを取り出して setattr でテスト用の設定値を注入してあげましょう

サンプルコードはこちら

github.com

もうちょっと詳しく

背景

PythonでDynamoDBを使った開発していればPynamoDBはとても便利なライブラリです。非常に書きやすいAPIでDynamoDBを読み書きできますし、手軽にテーブル自体もPynamoDBで作成することも可能です。 *1

しかしPynamoDBでテーブルを作成・管理してしまうと if table.exists() みたいな条件を書いて毎回判断させたり、キャパシティの調整のたびにアプリケーションを動かす必要が出てきてしまいます。 *2

そもそもテーブルの作成やキャパシティの管理はアプリケーションの責務ではないため、本番環境ではPynamoDBの責務をあくまで「アプリケーションからDynamoDBを使うクライアント」ととしての用途に限定し、テーブル自体は

  • Terraformで定義する
  • CDKで定義する
  • Serverless Frameworkを使う場合は serverless.yml の中で定義する

あたりがよく採用される方法かと思います。Partition Key/Sort Keyの設定やインデックスの定義などのほか、キャパシティはProvisionedかOn Demandか、Provisionedであればどれくらい積んでおくのか、といった設定を上記のいずれかのInfrastructure as Codeな方法で管理することでしょう。

すると当然、PynamoDBを使ったアプリケーションコードの中にはキャパシティに関する記述は残したくありません。アプリケーションを動かす際はそれでOKですが、テストを書くときに少し面倒なことになります。

たとえばテストにpytestを使う場合、fixture定義の中で 「テスト用のDynamoDBテーブルをlocalstackやDynamoDB Localのようなツールを使って作成し、テストが終われば削除する」 ような処理を書きたくなると思います。

PynamoDBでは、class Meta の中で host を定義することで、boto3で endpoint_urlを指定したのと同じ効果が得られます。

from pynamodb.models import Model

class Thread(Model):
    class Meta:
        table_name = 'Thread'
        # Specifies the region
        region = 'us-west-1'
        # Optional: Specify the hostname only if it needs to be changed from the default AWS setting
        host = 'http://localhost'
        # Specifies the write capacity
        write_capacity_units = 10
        # Specifies the read capacity
        read_capacity_units = 10
    forum_name = UnicodeAttribute(hash_key=True)

(公式Doc より抜粋 )

host = 'http://localhost' if ENV == "test" else None など書けば、テスト時と本番とで設定を打ち分けることは出来そうです。しかし、 アプリケーションのコードに「これはテスト実行かどうか」を判定する条件文が入るのは好ましくありません。

素直にboto3を使っていれば、上記スライド *3 の続きで解説されているのと同様に ddb = boto3.client('dynamodb', endpoint_url='http://localhost') のように作成したオブジェクトに切り替えることで、比較的簡単にアプリケーションからテストのためのロジックを追い出すことが可能でしょう。

しかしPynamoDBを使っている場合は、APIを実際に叩く部分はライブラリの中に隠蔽されてしまい、オブジェクトの切り替えによるテストが難しくなってしまいます。

また、テスト用のテーブルをどう作成するかという問題もあります。せっかくPynamoDBモデルを書いたのだから、できればうまく再利用したいものです。しかし、PynamoDBの Model.create_table() は、キャパシティに関する指定がないとエラーになります。

AttributeError: type object 'Meta' has no attribute 'read_capacity_units'

host の指定の時と同様、 billing_mode = "PAY_PER_REQUEST" if ENV =="test" else None などと書けば回避できそうですが *4 、そうするとまたアプリケーションに「テストのためのロジック」が混入してしまいます。

そこで、今回は

  • アプリケーションのコードにテストのためのロジックを入れない
  • PynamoDBのモデル定義をテスト用のテーブル作成に活用する

を同時に達成できるテストの書き方を考えてみました。

やってみる

サンプルのテーブルの定義はこんな感じです。(GitHubのほうのコードでは、リアリティを出すためにGSIとかも足しています)

from pynamodb.models import Model

class _UserModel(Model):
    class Meta:
        table_name = USER_TABLE_NAME

    uid: UnicodeAttribute = UnicodeAttribute(hash_key=True)
    name: UnicodeAttribute = UnicodeAttribute()
    group: UnicodeAttribute = UnicodeAttribute()

それっぽいテストコードを書きたいので、Repositoryパターンっぽく包んでみます。

class UserRepo:
    def __init__(self) -> None:
        self.model = _UserModel

    def add_user(self, uid: str, name: str, group: str) -> dict[str, str]:
        # 省略 実装はGitHubのサンプルを参照してください

    def get_user(self, uid: str) -> dict[str, str]:
        # 省略 実装はGitHubのサンプルを参照してください

    def get_users_in_group(self, group: str) -> list[dict[str, str]]:
        # 省略 実装はGitHubのサンプルを参照してください

この UserRepo の中の self.model を使って create_table しようとすると、billing_mode や read_capacity_units / write_capacity_unitsの設定がないのでAttributeErrorになります。

そこで、pytestのfixtureをこんな感じで書いてみます。

@pytest.fixture(scope="function")
def user_repo() -> Iterable[UserRepo]:

    repo: UserRepo = UserRepo()

    model_meta_class = getattr(repo.model, "Meta") # 1
    setattr(model_meta_class, "host", DDB_LOCAL_HOST) #2
    setattr(model_meta_class, "billing_mode", "PAY_PER_REQUEST") # 3
    setattr(model_meta_class, "table_name", "user_table_for_test") # 4

    repo.model.create_table(wait=True)

    yield repo

    # Delete table after running a test function
    repo.model.delete_table()

何をしているかというと、

  • # 1 でモデル定義の中のメタクラスを取り出し
  • # 2 でそのメタクラスにテスト用のhost(boto3でいうendpoint_url)を設定
  • # 3 でcreate_tableを通すためにbilling_modeを設定
  • # 4 でテスト用のテーブル名に名前を上書き

といった操作をしています。

テストコードはこんな感じで、

def test_user_repo(user_repo: UserRepo) -> None:
    alice = user_repo.add_user(uid="001", name="Alice", group="Red")
    assert alice == {"uid": "001", "name": "Alice", "group": "Red"}

    bob = user_repo.add_user(uid="002", name="Bob", group="Blue")
    chris = user_repo.add_user(uid="003", name="Chris", group="Blue")

    users_in_blue_group = user_repo.get_users_in_group(group="Blue")
    assert users_in_blue_group == [bob, chris]

実際にテストを実行してみると、

docker-compose up -d
docker-compose exec app pytest

結果:

f:id:TatchNicolas:20201210214630p:plain
テスト結果

pytestのfixtureでのテーブル作成が成功し、テストが通りました！

おまけ

また、以前のブログで「Serverless Framework +FastAPI」の開発環境を作った際に「せっかくserverless.ymlのなかでスキーマ定義してるのに、結局PynamoDBでテーブル作るためにモデル定義に余計なコードが混ざってしまうなあ...」という問題が残っていました。

tech.jxpress.net

この問題も、本記事と同じ方法でローカル開発用(≠テスト用)のテーブル作成スクリプトをアプリケーションコードとは別に切り出すことで解決できそうですね。コンテナイメージやLambdaには app/ 以下のコードだけ載せておけば良いので、「アプリケーションに余計なコードが混ざらない」を達成できます。

.
├── app
│   ├── __init__.py
│   ├── config.py
│   ├── main.py
│   └── repository.py
├── create_local_table.py
└── test
    └── test_repository.py

create_local_table.py スクリプトの中身はほとんどpytestの中身と同じです。

from os import environ

from app.repository import UserRepo


DDB_LOCAL_HOST = environ["DDB_LOCAL_HOST"]


if __name__ == "__main__":
    repo: UserRepo = UserRepo()

    model_meta_class = getattr(repo.model, "Meta")
    setattr(model_meta_class, "host", DDB_LOCAL_HOST)
    setattr(model_meta_class, "billing_mode", "PAY_PER_REQUEST")
    setattr(model_meta_class, "table_name", "user_table_for_dev")
    # or
    # setattr(model_meta_class, "read_capacity_units", 1)
    # setattr(model_meta_class, "write_capacity_units", 1)

    by_group_meta_class = getattr(repo.model.by_group, "Meta")
    setattr(by_group_meta_class, "host", DDB_LOCAL_HOST)

    repo.model.create_table(wait=True)

まとめ

Pythonのbuilt-inな関数である getattr / setattr を使うだけですが、当初のねらいであった「アプリケーションのコードにテストのためのロジックを入れない 」「PynamoDBのモデル定義をテスト用のテーブル作成に活用する」を達成することができました。

少しでも参考になれば幸いです。

最後に

JX通信社では、PythonやGoを使って「NewsDigest」の開発に参加してくれるインターン生を募集しています！ サーバレス、コンテナなど色々な技術スタックに触れられる環境なので、興味のある方は是非お声かけください！

www.wantedly.com

*1:https://pynamodb.readthedocs.io/en/latest/quickstart.html#creating-a-model

*2:RDSの場合、Read Replicaの数やインスタンスサイズをInfrastracture as codeなツールに(=インフラの責務)、テーブルのスキーマはDjangoやAlembicなどにマイグレーション管理させる(=アプリケーションの責務)パターンになると思いますが、DynamoDBの場合は「テーブル自体の作成・キャパシティ設定(=インフラの責務)」と「キーやインデックスといった設定(=アプリケーションの責務)」という切り分けになり、どのツールに何を任せるかという問題かと思います

*3:Node.jsかつS3の例ですがポイントは同じで、「テストのためのロジックが紛れ込んでいる」ことを問題にしているので例として引用しました

*4:read_capacity_unitsとwrite_capacity_unitsの両方を適当な数値に指定しても回避できます

© 2022 JX PRESS Corp.