public note

Python のユニットテストで decorator に patch をあてる

ユニットテストでデコレータを差し替える方法を学びました。

このような、print や sleep するデコレータを例にとります。

import time

def my_decorator(message="hi", sleep=0):
    def _my_decorator(func):
        def wrapper(*args, **kwargs):
            print("before")
            if sleep > 0:
                print(message)
                print(f"sleep {sleep} sec")
                time.sleep(sleep)
            result = func(*args, **kwargs)
            print("after")
            return result
        return wrapper
    return _my_decorator

このデコレータを下記の my_func 関数に設定します。

from decorator import my_decorator

@my_decorator(message="hello", sleep=60)
def my_func():
    print("my_func")

if __name__ == "__main__":
    my_func()

これを実行すると下記のように標準出力へ出力されます。

$ python sample.py
before
hello
sleep 60 sec
my_func
after

この my_func 関数に対するユニットテストでは、デコレータの部分をスキップしたくなります。そこでデコレータを差し替えてみます。

import sys
import pytest
from unittest.mock import patch

@pytest.fixture(scope="function")
def patched_my_func():
    if "sample" in sys.modules:
        del sys.modules["sample"]
    with patch("decorator.my_decorator", lambda message, sleep: lambda func: func):
        from sample import my_func
        yield my_func

def test_my_func(patched_my_func):
    patched_my_func()

del sys.modules["sample"] で、読み込み済みの sample モジュールを削除しています。この sample モジュールにはデコレータ付きの my_func が含まれるためです。 このテストコードでは無くても大丈夫ですが、別のテストがある場合は必要になるはずです。

また、デコレータとその対象の関数はあえて別のモジュールにしています。my_func を import する前にデコレータに対する patch を有効にするためです。

ユニットテストの実行結果は下記のようになります。標準出力結果から、デコレータを無にできていることがわかります。

$ pytest -sv 
(省略)
test_sample.py::test_my_func my_func   PASSED