pytestとmotoを利用してAWSサービスのmockを使ったテストをしてみる

k.uehara

2023.07.13

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

今回は、pytestとmotoを利用してAWSサービスのmockを使ったテストをしてみたいと思います。

はじめに

一応、pytestとmotoについて簡単に説明します。

pytestについて

pytestは、Pythonで単体テストを行うための高機能なフレームワークです。

本フレームワークはテストデータの作成や、異なる条件でのテストを行うための機能などが提供されています。

使いやすさと拡張性に重点を置いて設計されており、テストが失敗したときのエラーメッセージも分かりやすく、原因特定がしやすくなっています。

motoについて

motoはAWSの各種サービスをモックするためのPythonライブラリです。

AWSのサービスとのやり取りをmock化することで、実際のAWS環境には接続せずにテストを行うことができます。

SSM Parameter Storeを使ったテスト

まず最初に、SSM Parameter Storeをmockを使ってテストしてみたいと思います。

テスト①

テストにあたり、以下のファイルを用意します。

/
├ data
  └ parameter_list.tsv
├ sample
  ├ __init__.py
  └ ssm_sample.py
└ test
  ├ __init__.py
  ├ mock_ssm_ps.py
  └ test_ssm_sample.py

まずdataフォルダ配下のparameter_list.tsvですが、これはParameter Storeに登録するパラメータ値のファイルになります。

parameter_list.tsv

/test/param	String	my_param
/test/secure_param	SecureString	my_secure_param

実環境のAWSには、既にこの値が登録されていることを想定します。

sampleフォルダ配下のssm_sample.pyはテスト対象のファイルです。次のように、keyを引数に受け取り、Parameter Storeから値を取得する関数を持ちます。

ssm_sample.py

import boto3


def get_parameter(param_key, WithDecryption=True):
    ssm_client = boto3.client("ssm")
    
    response = ssm_client.get_parameters(
        Names=[
            param_key,
        ],
        WithDecryption=WithDecryption,
    )
    if len(response["Parameters"]) > 0:
        return response["Parameters"][0]["Value"]
    else:
        return None

実環境のAWSには既にParameter Storeにデータが入っていますが、mockを利用する場合、Parameter Storeのmockにデータを入れることから始める必要があります。

そこで、testフォルダ配下にmock_ssm_ps.pyを用意します。

mock_ssm_ps.py

import boto3

PARAM_DATA_FILE = "../data/parameter_list.tsv"


def prepare_ssm_parameters():
    ssm = boto3.client("ssm", "ap-northeast-1")

    with open(PARAM_DATA_FILE, encoding="utf-8", newline="") as f:
        for line in f:
            cols = [x.strip() for x in line.split("\t")]
            ssm.put_parameter(Name=cols[0], Value=cols[2], Type=cols[1])

prepare_ssm_parameters()関数では、data配下のparameter_list.tsvを読み取り、Parameter Storeにputする処理を記載しています。

これで準備ができたので、テストファイルであるtest_ssm_sample.pyを作成します。

test_ssm_sample.py

import pytest
from moto import mock_ssm

from test import mock_ssm_ps
from sample import ssm_sample


@pytest.fixture(autouse=True)
def setup_ssm_ps():
    mock = mock_ssm()
    mock.start()
    mock_ssm_ps.prepare_ssm_parameters()
    yield
    mock.stop()


def test_get_parameter():
    ret1 = ssm_sample.get_parameter("/test/param")
    ret2 = ssm_sample.get_parameter("/test/secure_param")

    assert [ret1, ret2] == ["my_param", "my_secure_param"]

ここではpytestのfixtureを用いています。

fixtureはテストの前処理（DBをセットアップする、mockを作成する等）を行うためのpytestの機能です。

fixture(autouse=True)とすることにより、全てのテスト関数の実行前にセットアップ用の関数が呼び出されるようにしています。

fixtureにはscopeという引数もあり、function, class, module, package, sessionを割り当てることができます。（例えば、テスト実行前にSMTPのコネクション確立の関数を1度だけ呼び出したい、というケースではscopeはmoduleにすると良さそうです）

今回は特に指定していないため、デフォルトはfunctionになっており、全てのテスト関数実行前に呼び出される形となります。

ちなみにautouse=Trueとせず、明示的にセットアップを利用したい関数を指定するには以下のように記載します。

import pytest


@pytest.fixture()
def setup_hoge():
    ...


def test_aaa(setup_hoge):
    ...

テストの実行

pytestの実行は以下で可能です。

$ pytest -vv test_ssm_sample.py

以下のようにテストがPASSしていたら成功です。

テスト②

fixtureの確認のため、ssm_sample.pyとtest_ssm_sampleにそれぞれ関数を追加してみます。

ssm_sample.py

import boto3


def get_parameter(param_key, WithDecryption=True):
    ssm_client = boto3.client("ssm")
    
    response = ssm_client.get_parameters(
        Names=[
            param_key,
        ],
        WithDecryption=WithDecryption,
    )
    if len(response["Parameters"]) > 0:
        return response["Parameters"][0]["Value"]
    else:
        return None


# 追加
def get_parameter_and_join_word(param_key, word="", WithDecryption=True):
    ssm_client = boto3.client("ssm")
    
    response = ssm_client.get_parameters(
        Names=[
            param_key,
        ],
        WithDecryption=WithDecryption,
    )
    if len(response["Parameters"]) > 0:
        return response["Parameters"][0]["Value"] + word
    else:
        return None

ssm_sample.pyには、引数として受け取ったキーからssmのパラメータの値を取得し、同じく引数として受け取ったwordをパラメータの値の末尾に追加するという関数を追加してみました。

この関数をテストするテスト関数をtest_ssm_sample.pyに追加します。

test_ssm_sample.py

import pytest
from moto import mock_ssm

from test import mock_ssm_ps
from sample import ssm_sample


@pytest.fixture(autouse=True)
def setup_ssm_ps():
    mock = mock_ssm()
    mock.start()
    mock_ssm_ps.prepare_ssm_parameters()
    yield
    mock.stop()


def test_get_parameter():
    ret1 = ssm_sample.get_parameter("/test/param")
    ret2 = ssm_sample.get_parameter("/test/secure_param")

    assert [ret1, ret2] == ["my_param", "my_secure_param"]


# 追加
def test_get_parameter_and_join_word():
    ret = ssm_sample.get_parameter_and_join_word("/test/param", "_hoge")

    assert ret == "my_param_hoge"

テストの実行

先程と同様にpytestを実行します。

$ pytest -vv test_ssm_sample.py

以下のように、2つのテストがPASSしたら成功です。

DynamoDBを使ったテスト

先程のSSMに加え、DynamoDBのテストもしてみます。

テスト①

新規に追加するファイルは以下の通りです。

/
├ data
  ├ parameter_list.tsv
  ├ ddb_foo_table_data.json　　★新規
  └ ddb_foo_table_def.json  ★新規
├ sample
  ├ __init__.py
  ├ ssm_sample.py
  └ ddb_sample.py           ★新規
└ test
  ├ __init__.py
  ├ mock_ssm_ps.py
  ├ mock_ddb.py             ★新規
  ├ test_ssm_sample.py
  └ test_ddb_sample.py      ★新規

せっかくなのでSSMとDynamoDBを同時に使ってみたいので、Parameter Storeの設定値を記載しているparameter_list.tsvにDynamoDBのテーブル名をパラメータの値として1行追記します。

parameter_list.tsv

/test/param	String	my_param
/test/secure_param	SecureString	my_secure_param
/test/ddb_table_name	String	foo

次にdataフォルダ配下にDynamoDBのテーブル定義を記載したddb_foo_table_def.jsonと、テーブル内のデータを記載したddb_foo_table_data.jsonを作成します。

ddb_foo_table_def.json

{
    "AttributeDefinitions": [
        {
            "AttributeName": "dataId",
            "AttributeType": "S"
        }
    ],
    "KeySchema": [
        {
            "AttributeName": "dataId",
            "KeyType": "HASH"
        }
    ],
    "BillingMode": "PAY_PER_REQUEST"
}

ddb_foo_table_data.json

[
    {
        "PutRequest": {
            "Item": {
                "dataId": "1",
                "data": "data_01"
            }
        }
    },
    {
        "PutRequest": {
            "Item": {
                "dataId": "2",
                "data": "data_02"
            }
        }
    }
]

sampleフォルダ配下のddb_sample.pyはテスト対象のファイルです。

data_idを引数に受け取り、Parameter Storeから取得したDynamoDBのテーブルから値を取得するget_item_from_data_id()関数をテスト対象の関数としたいと思います。

ddb_sample.py

import boto3

from sample import ssm_sample


def _connect_dynamodb(table_name):
    dynamodb = boto3.resource("dynamodb")
    return dynamodb.Table(table_name)


def _query_key(table_name, key):
    table = _connect_dynamodb(table_name)
    result = table.get_item(Key=key)
    result_item = result.get("Item")
    return result_item


def get_item_from_data_id(data_id):
    table_name = ssm_sample.get_parameter("/test/ddb_table_name")
    result = _query_key(table_name, {"dataId": data_id})
    return result

Parameter Storeと同様、実環境のAWSには既にDynamoDBにデータが入っていますが、mockを利用する場合、DynamoDBのmockにデータを入れることから始める必要があります。

そこで、testフォルダ配下にmock_ddb.pyを用意します。

mock_ddb.py

import json

import boto3

TABLE_DEF_FILE = "../data/ddb_foo_table_def.json"
TABLE_DATA_FILE = "../data/ddb_foo_table_data.json"


def prepare_foo_table():
    table_name = "foo"
    dynamodb = boto3.resource("dynamodb")

    # テーブルの作成
    with open(TABLE_DEF_FILE) as f:
        foo_table_def = json.load(f)
    dynamodb.create_table(
        TableName=table_name,
        **foo_table_def,
    )

    # テーブルにデータを投入
    with open(TABLE_DATA_FILE) as f:
        data_str = f.read()
    data_items = json.loads(data_str)
    table = dynamodb.Table(table_name)
    for item in data_items:
        table.put_item(Item=item["PutRequest"]["Item"])

prepare_foo_table()関数では、data配下のテーブル定義とデータを読み取り、DynamoDBにputする処理を記載しています。

これで準備ができたので、テストファイルであるtest_ddb_sample.pyを作成します。

test_ddb_sample.py

import pytest
from moto import mock_dynamodb, mock_ssm

from test import mock_ddb, mock_ssm_ps
from sample import ddb_sample


@pytest.fixture(autouse=True)
def setup_ssm_ps():
    mock = mock_ssm()
    mock.start()
    mock_ssm_ps.prepare_ssm_parameters()
    yield
    mock.stop()


@pytest.fixture(autouse=True)
def setup_ddb():
    mock = mock_dynamodb()
    mock.start()
    mock_ddb.prepare_foo_table()
    yield
    mock.stop()


def test_get_item_from_data_id():
    ret = ddb_sample.get_item_from_data_id("1")["data"]
    assert ret == "data_01"

ここでは、fixtureでParameter StoreとDynamoDBの両方のセットアップ用関数を用意しています。

テストの実行

pytestを実行します。

$ pytest -vv test_ddb_sample.py

以下のように、テストがPASSしたら成功です。

テスト②

fixtureを使わない、別の書き方もご紹介します。

例えばDynamoDBをmock化したい場合、mock化したいテスト関数の先頭に@mock_dynamodbを記述することでも実現可能です。

test_ddb_sample.py

import pytest
from moto import mock_dynamodb, mock_ssm

from test import mock_ddb, mock_ssm_ps
from sample import ddb_sample


@pytest.fixture(autouse=True)
def setup_ssm_ps():
    mock = mock_ssm()
    mock.start()
    mock_ssm_ps.prepare_ssm_parameters()
    yield
    mock.stop()


@mock_dynamodb
def test_get_item_from_data_id():
    mock_ddb.prepare_foo_table()
    
    ret = ddb_sample.get_item_from_data_id("1")["data"]
    assert ret == "data_01"

ただ、Parameter StoreやDynamoDBは様々な関数から呼ばれるケースも多いと思いますので、fixtureを使いセットアップ用の関数を用意しておくと便利です。

fixtureを利用しない場合、上記の例だとDynamoDBを使う全てのテスト関数で@mock_dynamodbの記載とmock_ddb.prepare_foo_table()の記述が必要となります。

テストの実行

pytestを実行します。

$ pytest -vv test_ddb_sample.py

以下のように、テストがPASSしたら成功です。

最後に

今回は、pytestとmotoを利用してAWSサービスのmockを使ったテストをしてみました。

参考になりましたら幸いです。

参考文献

AWS

関連記事

Amazon BedrockとSlackで生成AIチャットボットアプリを作る (その1：ローカルPC＋Socket Modeで動かす)

青柳英明

2024.05.06

Data Wrangler Extension for Visual Studio Code がGAになったよ

大谷 一輝

2024.05.01

Pythonでリトライ処理を簡単に追加できる「tenacity」を使ってみる

アシ

2024.04.30

話題の記事

Pythonで日時の文字列にタイムゾーンをつけるときの注意点

kobayashi.m

2024.04.30