boto3をpytestのmonkeypatchでstubする

Sep 24, 2023

背景

業務では boto3 を使って取得したデータをデータベースに格納するスクリプトを書きました。
テストを書くために boto3 を stub する方法を調べてると、stubber というものがありました。
boto3 を stub する方法について調査すると、ネットには stubber と unittest の組み合わせを使った方法が多く紹介されていました。
しかし stubber と pytest の mock を組み合わせたは記事が少ないように感じました。

今回は boto3 を pytest の monkeypatch で stub する方法を備忘録として残します。

環境構築

$ python3 -m venv venv
$ source venv/bin/activate

必要なモジュールをインストールします。

$ pip install boto3 pytest boto3-stubs['iam']

boto3 を実行できるようにするため環境変数を設定します。

export AWS_ACCESS_KEY_ID=AKxxxxxxxxxxxxx
export AWS_SECRET_ACCESS_KEY=wJaxxxxxxxxxxxxxxxxxxx
export AWS_DEFAULT_REGION=ap-northeast-1

exampleクラス

今回は Example クラスにある __get_client を patch します。

src/example.py

from typing import Any
import boto3
from mypy_boto3_iam.client import IAMClient

class Example:
  def __init__(self) -> None:
    self.__client = self.__get_client()

  def __get_client(self) -> IAMClient:
    return boto3.client("iam")

  def get_roles(self) -> Any:
    return self.__client.list_roles()

エントリーポイントは main.py です

src/main.py

import json
from datetime import date, datetime
from example import Example

def json_serial(obj):
    if isinstance(obj, (datetime, date)):
        return obj.isoformat()
    raise TypeError ("Type %s not serializable" % type(obj))


ex = Example()

print(json.dumps(ex.get_roles(), default=json_serial))

• [Python] dateやdatetimeをjson.dumpでエラーなく出力する

コードを実行するとロール一覧を取得できます。

$ python src/main.py | jq
{
  "Roles": [
    {
      "Path": "/",
      "RoleName": "AdminRole",
      ...

pytestでテストを書く

Example クラスにある get_roles メソッドをテストしたいと思います。

stubber で iam client の stub を作成する

Stubber Reference にある Example を参考にしてコード書いてみます。

tests/test_example_1.py

import json
import botocore.session
from botocore.stub import Stubber, ANY

from src.example import Example

iam = botocore.session.get_session().create_client('iam')
stubber = Stubber(iam)


response = """
{
  "Roles": [
    {
      "Path": "/",
      "RoleName": "TestRole",
      "RoleId": "AAAABBBCCDDDADAAAAAAA",
      "Arn": "arn:aws:iam::123456789d:role/TestRole",
      "CreateDate": "2021-02-20T23:16:13+00:00",
      "Description": "",
      "MaxSessionDuration": 3600
    }
  ]
}
"""

expected_params = {}
stubber.add_response("list_roles", json.loads(response), expected_params)

def test_example(monkeypatch):
  stubber.activate()
  monkeypatch.setattr(Example, "_Example__get_client", lambda _: iam)

  res = Example().get_roles()
  assert res.get("Roles")[0].get("RoleName") == "TestRole"
  stubber.deactivate()

monkeypatch には fixture がついているので、テストの関数やメソッドには引数に monkeypatch を取ることができます。 https://github.com/pytest-dev/pytest/blob/d2b214220f63e1fc90120495d600893cfba6219f/src/_pytest/monkeypatch.py#L29-L30 info: pytest fixture 一覧

テストは問題なさそうです。

$ pytest -s tests
==================================== test session starts ====================================
platform darwin -- Python 3.11.4, pytest-7.4.2, pluggy-1.3.0
rootdir: /Users/xxxxxx/blog-code/2023/09/boto3-monkeypatch
collected 1 item

tests/test_example.py .

===================================== 1 passed in 0.33s ====================================

テストコードには monkeypatch.setattr を使って Example クラスのプライベートメソッドを置き換えました。 monkeypatch.setattr(Example, "_Example__get_client", lambda _: iam)

プライベートメソッド名

プライベートメソッドの名前は、_classname_private-methodname になります。
これは組み込み関数の dir を使うと確認できます。

src/main.py

import json
from example import Example

ex = Example()

print(dir(Example))
# print(ex.get_roles())

__get_client の名前は、_Example__get_client であることがわかります。

$ python src/main.py
['_Example__get_client', ...

__spam (先頭に二個以上の下線文字、末尾に一個以下の下線文字) という形式の識別子は、 _classname__spam へとテキスト置換されるようになりました。

引用：https://docs.python.org/ja/3/tutorial/classes.html#private-variables

stub の作成を fixture にする

stubber の箇所を fixture にして再利用しやすいようにしてみます。

tests/test_example_2.py

import json
import botocore.session
from botocore.stub import Stubber, ANY
import pytest
from mypy_boto3_iam.client import IAMClient

from src.example import Example

@pytest.fixture
def iam_client() -> tuple[IAMClient, Stubber]:

  response = """
  {
    "Roles": [
      {
        "Path": "/",
        "RoleName": "TestRole",
        "RoleId": "AAAABBBCCDDDADAAAAAAA",
        "Arn": "arn:aws:iam::123456789d:role/TestRole",
        "CreateDate": "2021-02-20T23:16:13+00:00",
        "Description": "",
        "MaxSessionDuration": 3600
      }
    ]
  }
  """

  iam = botocore.session.get_session().create_client('iam')
  stubber = Stubber(iam)

  expected_params = {}
  stubber.add_response("list_roles", json.loads(response), expected_params)

  return iam, stubber


def test_example(monkeypatch, iam_client):
  iam, stubber = iam_client

  stubber.activate()
  monkeypatch.setattr(Example, "_Example__get_client", lambda _: iam)

  res = Example().get_roles()
  assert res.get("Roles")[0].get("RoleName") == "TestRole"
  stubber.deactivate()

無事にテストも通りました。これで stub する箇所を再利用できるようになりました。

$ pytest -s tests
==================================== test session starts ====================================
platform darwin -- Python 3.11.4, pytest-7.4.2, pluggy-1.3.0
rootdir: /Users/xxxx/blog-code/2023/09/boto3-monkeypatch
collected 2 items

tests/test_example_1.py .
tests/test_example_2.py .

===================================== 2 passed in 0.22s ====================================

リファクタリングする

テスト用のクラスを作成したい場合もあると思います。
pytest の fixtureの説明で、“Arrange”、“Act”、“Assert”、“Cleanup” の話があるので、これに沿ってリファクタしてみます。

tests/test_example_3.py

import json
import botocore.session
from botocore.stub import Stubber, ANY
import pytest
from mypy_boto3_iam.client import IAMClient

from src.example import Example

@pytest.fixture
def iam_client() -> tuple[IAMClient, Stubber]:

  response = """
  {
    "Roles": [
      {
        "Path": "/",
        "RoleName": "TestRole",
        "RoleId": "AAAABBBCCDDDADAAAAAAA",
        "Arn": "arn:aws:iam::123456789d:role/TestRole",
        "CreateDate": "2021-02-20T23:16:13+00:00",
        "Description": "",
        "MaxSessionDuration": 3600
      }
    ]
  }
  """

  iam = botocore.session.get_session().create_client('iam')
  stubber = Stubber(iam)

  expected_params = {}
  stubber.add_response("list_roles", json.loads(response), expected_params)

  return iam, stubber


class TestExample():

  def test_example(self, monkeypatch, iam_client):
    # ============
    #  Arrange
    # ============
    iam, stubber = iam_client
    stubber.activate()
    monkeypatch.setattr(Example, "_Example__get_client", lambda _: iam)

    # ============
    #  Act
    # ============
    res = Example().get_roles()

    # ============
    #  Assert
    # ============
    assert res.get("Roles")[0].get("RoleName") == "TestRole"

    # ============
    #  Clearnup
    # ============
    stubber.deactivate()

記述量が増えてしまいましたが、まとまりがわかりやすいためテストを追加しやすくなったと思います。

$ pytest -s tests
==================================== test session starts ====================================
platform darwin -- Python 3.11.4, pytest-7.4.2, pluggy-1.3.0
rootdir: /Users/xxxx/blog-code/2023/09/boto3-monkeypatch
collected 3 items

tests/test_example_1.py .
tests/test_example_2.py .
tests/test_example_3.py .

===================================== 3 passed in 0.20s ====================================

最後に

今回使用したコードは、下のリンクに置いてあります。 https://github.com/kntks/blog-code/tree/main/2023/09/boto3-monkeypatch

pytest の monkeypatch と boto3 の stubber の組み合わせに焦点を当てて記事を書いたため、全体的に説明不足なところもあると思いますが、テストコードを書くときの参考になれば幸いです。

Tags:

• Python
• AWS

【Go】slog を触ってみる

pydanticを触ってみる