肉球でキーボード

MLエンジニアの技術ブログです

pymock でモック化した関数の特定の引数の値を取得する方法

pymock でモック化した際に、モックが呼び出された時の特定の引数の値を取得する機会があったのでやり方を紹介します。

やりたいこと

以下のような、外部で定義した関数(例: calc_sum)を呼び出す処理を行うメソッド(例: process_{one, two})をテストする際に、メソッド中の関数 (calc_sum) をモック化した場合、モックが呼び出された時の各引数の値をテスト関数で取得する状況を考えます。

from typing import List


def calc_sum(arg1: List[float], arg2: List[float]) -> float:
    """Calculate sum of two argments"""
    return sum(arg1) + sum(arg2)


def process_one(arg1: List[float], arg2: List[float]) -> float:
    """Use calc_sum function ones in process"""
    sum_of_args = calc_sum(arg1=arg1, arg2=arg2)

    return sum_of_args


def process_two(arg1: List[float], arg2: List[float], arg3:List[float], arg4:List[float]) -> (float, float):
    """Use calc_sum function twice in process"""
    sum_of_args = calc_sum(arg1=arg1, arg2=arg2)
    sum_of_args2 = calc_sum(arg1=arg3, arg2=arg4)

    return sum_of_args, sum_of_args2

1度だけモックが呼ばれた場合

1 度だけモックが呼ばれた場合の引数の取得場合は call_argsメソッドを用います。

call_args 公式ドキュメント

注意して欲しいのは、call_args メソッドは引数の辞書型を返すのではなく、unittest.mock._Call オブジェクトを返す点です。

_Call クラスの docstring を見てみると、以下のようなタプル型となっています。 https://github.com/python/cpython/blob/324b93295fd81133d2dd597c3e3a0333f6c9d95b/Lib/unittest/mock.py#L2381-L2399

class _Call(tuple):
    """
    A tuple for holding the results of a call to a mock, either in the form
    `(args, kwargs)` or `(name, args, kwargs)`.
    If args or kwargs are empty then a call tuple will compare equal to
    a tuple without those values. This makes comparisons less verbose::
        _Call(('name', (), {})) == ('name',)
        _Call(('name', (1,), {})) == ('name', (1,))
        _Call(((), {'a': 'b'})) == ({'a': 'b'},)
    The `_Call` object provides a useful shortcut for comparing with call::
        _Call(((1, 2), {'a': 3})) == call(1, 2, a=3)
        _Call(('foo', (1, 2), {'a': 3})) == call.foo(1, 2, a=3)
    If the _Call has no name then it will match any name.
    """

_Call オブジェクトはタプル型なので、要素にアクセスすることでモックが呼び出された際の引数の値を取得することができます。

from unittest.mock import MagicMock, call, patch

from src.func import process_one, process_two


@patch("src.func.calc_sum")
def test_process_one(mocked_calc_sum: MagicMock):
    actual = process_one(arg1=[0.1, 0.2], arg2=[0.3, 0.4])
    actual_args = mocked_calc_sum.call_args
    print(actual_args) # call(arg1=[0.1, 0.2], arg2=[0.3, 0.4])
    print(actual_args[0]) # ()
    print(actual_args[1]) # {'arg1': [0.1, 0.2], 'arg2': [0.3, 0.4]}

複数回モックが呼ばれた場合

call_args メソッドはモックが最後に呼び出された引数を取得するため、複数回呼ばれた時の引数全てを取得することはできません。
複数回呼ばれた時の引数を全て取得する場合は call_args_list メソッドを使用します。
call_args_list 公式ドキュメント

ここで注意するのは、call_args_list によって取得できるのは _Call オブジェクトのリストであるという点です。

call_args_list の実装を見ると、Call オブジェクトをリストに追加していますね。

https://github.com/python/cpython/blob/324b93295fd81133d2dd597c3e3a0333f6c9d95b/Lib/unittest/mock.py#L1099-L1108

def _increment_mock_call(self, /, *args, **kwargs):
        self.called = True
        self.call_count += 1

        # handle call_args
        # needs to be set here so assertions on call arguments pass before
        # execution in the case of awaited calls
        _call = _Call((args, kwargs), two=True)
        self.call_args = _call
        self.call_args_list.append(_call)

各呼び出し引数を取得する場合、リストから対象の呼び出し引数の_Call オブジェクトを取得し、引数の要素にアクセスします。

@patch("src.func.calc_sum")
def test_process_two(mocked_calc_sum: MagicMock):
    actual = process_two(arg1=[0.1, 0.2], arg2=[0.3, 0.4], arg3=[0.5, 0.6], arg4=[0.7, 0.8])
    actual_args = mocked_calc_sum.call_args_list
    print(actual_args) # [call(arg1=[0.1, 0.2], arg2=[0.3, 0.4]), call(arg1=[0.5, 0.6], arg2=[0.7, 0.8])]
    print(actual_args[0][1]) # {'arg1': [0.1, 0.2], 'arg2': [0.3, 0.4]}
    print(actual_args[1][1]) # {'arg1': [0.5, 0.6], 'arg2': [0.7, 0.8]}
 

参考

unittest.mock --- モックオブジェクトライブラリ — Python 3.10.4 ドキュメント

Pythonで、モックに差し替えたメソッドが呼ばれた回数や呼ばれた時の引数を検証する - メモ的な思考的な