Python でリトライ処理のテストを書く (pytest)

Python でリトライ処理を書いたのでメモ。

テスト対象のコード

target.py

class ApiConnectionError(Exception):
    pass


class Target:

    def get(self, key: str) -> dict[str, str]:
        """api からデータを取得する処理

        Args:
            key (str): API パスパラメーター

        Returns:
            dict[str, str]: API レスポンス

        Note:
            API 通信に失敗すると
            ApiConnectionError を送出します
        """        
        # key を使って api からデータを取得する処理が書かれている
        return {"get": "data"}
    
    def retry_get(self, key: str) -> dict[str, str]:
        """1回だけ self.get をリトライする関数

        Args:
            key (str): API パスパラメーター

        Returns:
            dict[str, str]: API レスポンス
        """        
        try:
            return self.get(key)
        except ApiConnectionError:
            return self.get(key)

プロジェクト構成は以下のようになっています。

❯ tree --gitignore
.
├── target
│   ├── __init__.py
│   └── target.py
└── tests
    ├── __init__.py
    └── test_target.py

pytest でのテスト

retry_get は get を1回だけリトライするので、ApiConnectionError が発生したときに get が本当に2回実行されているかをテストします。

test_target.py

from unittest import mock

import pytest

class TestTarget:

    def test_retry_execute(self):
        """get が 2回呼ばれることをテストする
        """        
        from target.target import Target, ApiConnectionError

        t = Target()
        # get が ApiConnectionError を送出するようにモック
        t.get = mock.Mock(side_effect=ApiConnectionError)

        with pytest.raises(ApiConnectionError):
            t.retry_get("test")

        assert t.get.call_count == 2

t.get に ApiConnectionError を送出してもらいたいので、 Mock(side_effect=ApiConnectionError) のようにモックします。
side_effect については unittest.mock --- 入門 を参照してください。

次に t.get が例外を送出するようにモックしたので、t.retry_get も例外を送出するようになります。
ここで、pytest.raises() をコンテキストマネージャとして利用して t.retry_get を実行します。

最後に t.get (Mock) が呼ばれた回数を call_count 属性を使って検証します。

pytest を実行し成功すれば実際に t.get が2回実行されたことが検証できます！

おわりに

もっと良いやり方などあるという方がいましたらぜひ教えてください！

ここまで読んでくれてありがとうございました。