はじめに

本記事は Progaku Advent Calendar 2024 21日目の記事になります。

プログラムにはuuidをはじめとしてそのタイミングでランダムに生成される値があります。
そういった処理が含まれる単体テストでは何らかの固定値にモックしてテストを通す手法が考えられます。
また、ランダムな値であったとしても、生成された結果に応じて後続処理で条件分岐が入る場合や、生成のタイミングで失敗した後に処理が存在する場合もあります。
そういった場合は、それぞれテストを切って個別に固定値をモックして……としなければいけないのでしょうか？
私ならば、pytestを使っているのならば@pytest.mark.parametrizeを使ってまとめたいと考えます。

それを叶えるための柔軟な返り値のモック化について、本記事ではunittest.mockを用いた解決策を深掘ります。

環境構築

私が慣れているという理由でpytestを使います。

python3 -m venv .venv
./.venv/bin/pip install pytest

# 入る
source .venv/bin/activate
# test
pytest
# 出る
deactivate

今回検証するコード

from uuid import uuid4


class RetryOverException(Exception):
    pass


already_uuid = {}
RETRY_COUNT = 5


def _create(id: int, code: str) -> tuple[int, str, str]:
    return (id, code, '仮名')


def _get_uuid() -> str:
    cnt = 0
    while cnt < RETRY_COUNT:
        u = uuid4().hex
        if u not in already_uuid:
            return u

    raise RetryOverException()


def _get_user_data(id: int) -> tuple[int, str, str] | None:
    id_map = {
        1: (1, 'alice', 'アリス'),
        2: (2, 'bob', 'ボブ'),
        3: (3, 'john', 'ジョン'),
    }
    return id_map.get(id)


def get_or_create_user_data(id: int) -> tuple[int, str, str]:
    user_data = _get_user_data(id)
    try:
        if not user_data:
            code = _get_uuid()
            user_data = _create(id, code)
    except RetryOverException:
        raise Exception('リトライ回数上限に達しました')

    return user_data


def bulk_create(n: int) -> list[tuple[int, str, str]]:
    res = []
    try:
        for id in range(n):
            code = _get_uuid()
            res.append(_create(id, code))
    except RetryOverException:
        raise Exception('リトライ回数上限に達しました')

    return res

単体テスト

固定値の返却

まずは返り値を固定値にモックしたいケースです。
return_valueを用いて、生成するuuidを固定値にしています。

@mock.patch('main._get_uuid', return_value='uuid')
def test_get_or_create_user_data(mock_get_uuid):
    res = get_or_create_user_data(10)
    assert res == (10, 'uuid', '仮名')

エラーを発生させる

次は関数の実行結果がエラーになってしまった場合です。

side_effectを用いて、関数が指定Exceptionをraiseするようにしています。
本来であればこのテストは uuid.uuid4() が固定値を返すようにモックするのがいいのですが、今回は記事の都合で関数自体をモックしています。

@mock.patch('main._get_uuid', side_effect=RetryOverException)
def test_failed_issue_uuid_retry_over(mock_get_uuid):
    with pytest.raises(Exception) as e:
        get_or_create_user_data(10)
    assert str(e.value) == 'リトライ回数上限に達しました'

条件に応じて返り値を変える

side_effectに関数を渡すことで、モック対象の処理を書き換えることができます。
これを用いて条件分岐に応じた固定値を返すようにしました。
@pytest.mark.parametrize でテスト関数をまとめることができるようになりました。(下記ケースではreturn_valueの値を条件分岐であらかじめ用意しておく手法でもできる)

@pytest.mark.parametrize(
    ['id', 'expected'],
    [
        (0, (0, 'full_tuning', 'フルチューニング')),
        (999, (20001, 'last_order', '打ち止め')),
    ]
)
def test_branch_get_or_create_user_data(id, expected):
    def mock_get_user_data(arg_id: int):
        if arg_id == 0:
            return (0, 'full_tuning', 'フルチューニング')
        else:
            return None

    with mock.patch(
        'main._get_user_data', side_effect=mock_get_user_data
    ), mock.patch(
        'main._create', return_value=(20001, 'last_order', '打ち止め')
    ):
        res = get_or_create_user_data(id)
        assert res == expected

イテラブルに返り値を決める

side_effectを見ていると、イテレータの値を渡すことで、呼ばれるごとに返り値が変わっている例があります。

>> mock = Mock()
>> mock.side_effect = [3, 2, 1]
>> mock(), mock(), mock()
(3, 2, 1)

つまりはyieldを用いれば、呼ばれる回数を気にすることなく複数回呼ばれるけどそのたびに任意のルールを持った別の値を返して欲しいニーズを叶えることができます。
とはいえ、テストなのに呼ばれる回数を決め打たないなんてことはないと思うので、固定値のイテレーターを入れる方が健全だと思います。

def test_bulk_create():
    def mock_get_uuid():
        counter = 0
        while True:
            yield 'id_{:0>4}'.format(str(counter))
            counter += 1

    with mock.patch('main._get_uuid', side_effect=mock_get_uuid()):
        actual = bulk_create(10)

    assert actual == [
        (i, 'id_{:0>4}'.format(str(i)), '仮名')
        for i in range(10)
    ]

テスト

$ pytest
============================= test session starts ==============================
rootdir: /path/to/test
collected 5 items                                                              

test_main.py .....                                                       [100%]

============================== 5 passed in 0.02s ===============================

おわりに

unittest.mockはときどきパス指定とPythonのimport周りの挙動に悩まされながらも愛用させていただいています。
ただ、普段だと return_value で固定値、 side_effect で任意エラーの発生にしか使っておらず、 side_effect がここまで柔軟に使えるものだったとは、つい最近まで知りませんでした。
ただあくまでテストコードなので、柔軟に作りすぎると今度は開発者が結果の正しさを確認するのにテスト側のコードを追う手間が発生しますので、ほどほどにしましょう。

以上、Progaku Advent Calendar 2024 21日目でした。