Pythonで、unittest.mock.patchを使ってデコレータを差し替える

同僚と話している中で、unittest.mock.patchを使ったデコレータの差し替えに関する話題が出ました。

そういえばデコレータは差し替えたことがなかったため、試してみたことをメモします。

なお、「テストファイル群に、デコレータを差し替える/差し替えないものが混在している場合」で使った方法は強引な気がします。そのため、もしより良い方法をご存じであればご指摘ください。

 
目次

 

環境

  • Python 3.6.1
  • pytest 3.0.7
    • テストランナーとして使用

 

用意したデコレータとプロダクションコード

デコレータとは、関数に処理を追加するためのシンタックスシュガー(糖衣構文)です。

以下が参考になりますので、デコレータの詳細は省略します。

 
今回用意するデコレータは、

  • 引数なしのデコレータ (@countup, @countdown)
  • 位置引数ありのデコレータ (@add)
  • 位置引数とキーワード引数ありのデコレータ (@calculate)

です。

また、functools.wrapsを使っていますが、使う意味や使い方などは以下が参考になりました。
[python]デコレータでfunctools.wrap()を使う - logging.info(self)

 
デコレータのソースコードはこんな感じです。

deco/my_decorator.py

from functools import wraps

def countup(function):
    @wraps(function)
    def wrapper(*args, **kwargs):
        result = function(*args, **kwargs)
        return result + 1
    return wrapper


def countdown(function):
    @wraps(function)
    def wrapper(*args, **kwargs):
        result = function(*args, **kwargs)
        return result - 1
    return wrapper


def add(decorator_arg):
    """デコレータadd(@addとして使う)と、デコレータの引数decorator_arg"""

    def decorate(function):
        """デコレート対象の関数を引数fuctionとして受け取る、関数"""

        @wraps(function)
        def wrapped(*args, **kwargs):
            """デコレート対象の関数の引数をargsやkwargsとして受け取る、関数"""
            # デコレート対象の関数を実行し、結果をresultに入れる
            result = function(*args, **kwargs)
            # resultに対し、デコレータの引数を加算して戻す
            return result + decorator_arg

        # decorate関数は、デコレート対象の関数をラップするwrapped関数を返す
        return wrapped

    # add関数は、decorate関数を返す
    return decorate


def calculate(*decorator_args, **decorator_kwargs):
    u""""""
    def decorate(function):
        @wraps(function)
        def wrapped(*args, **kwargs):
            result = function(*args, **kwargs)
            # 可変長引数で与えられた数を合計する
            summary = sum(decorator_args)
            # キーワード可変長引数に減算指定がある場合は減算、それ以外は加算
            if decorator_kwargs.get('is_decrement'):
                return result - summary
            else:
                return result + summary
        return wrapped
    return decorate

 
これらのデコレータは、以下のようにしてプロダクションコードで使います。

target.py

from deco.my_decorator import countup, countdown, add, calculate

class Target:
    def __init__(self, value=0):
        self.value = value

    @countup
    def execute_count_up(self):
        return self.value

    @countdown
    def execute_count_down(self):
        return self.value

    @add(2)
    def execute_add(self):
        return self.value

    @calculate(1, 2, 3)
    def execute_calculate_increment(self):
        return self.value

    @calculate(1, 2, 3, is_decrement=True)
    def execute_calculate_decrement(self):
        return self.value

 
動作はテストコードで説明すると

test_target.py

from target import Target

class TestCountUpWithoutPatch:
    def test_count_up_decorator(self):
        sut = Target()
        actual = sut.execute_count_up()
        assert actual == 1

    def test_count_down_decorator(self):
        sut = Target()
        actual = sut.execute_count_down()
        assert actual == -1

    def test_add_decorator(self):
        sut = Target()
        actual = sut.execute_add()
        assert actual == 2

    def test_increment_decorator(self):
        sut = Target()
        actual = sut.execute_calculate_increment()
        assert actual == 6

    def test_decrement_decorator(self):
        sut = Target()
        actual = sut.execute_calculate_decrement()
        assert actual == -6

 
であり、テストはすべてパスします。

$ pytest test_target.py 
...

test_target.py .....

5 passed in 0.05 seconds

 

引数なしのデコレータを差し替える

プロダクションコード

@countup
def execute_count_up(self):
    return self.value

について、

  • execute_count_up()メソッドの挙動を確認したい
  • デコレータ@countupは動作させたくない

という条件を満たすテストを書きたいとします。

 
その方法を調べたところ、以下の記事がありました。
Patching decorators | Python Mock Gotchas - Alex Marandon

I hope this silly example convinces you that decorators are applied when the class is created, which happens straight away when we load the module containing the class.

To deal with this you first have to make sure that your decorator is defined in a module separate from the class, otherwise you’ll never get a chance to replace it with a mock before the class definition calls it. Then you need to write your test code so that it patches the decorator before it gets applied to the methods of your class:

とのことです。

そのため、プロダクションコードのimport前に、デコレータを差し替えれば良さそうです。

 
なお、上記の記事ではMockオブジェクトのstart()を使っていましたが、stop()するのがめんどうなので、今回はpatchをwith構文を使うことにします。

 
では、テストコードでデコレータを差し替えてみます。

test_global_patch_count_up.py

from unittest.mock import patch

# withを使ったパッチ
# デコレータはimportした時に確定するため、importだけをwithの中に入れる
with patch('deco.my_decorator.countup', lambda function: function):
    from target import Target

class TestCountUp:
    def test_count_up_decorator(self):
        sut = Target()
        actual = sut.execute_count_up()
        assert actual == 0

デコレータで何も処理しないよう、patchを使って差し替えました。

そのため、デコレータでは1を加算しなくなり、プロダクションコードからは1ではなく0が返ってくるはずです。

 
テストコードを実行してみます。

$ pytest test_global_patch_count_up.py 
...
test_global_patch_count_up.py .

...
1 passed in 0.08 seconds

想定通りの動きとなり、テストをパスしました。

 

引数ありのデコレータを差し替える

次は、引数ありのデコレータを差し替えるケースを考えます。

今回は、位置引数を持つデコレータ@addを差し替えます。

引数なしと同じように

test_global_patch_add.py

from unittest.mock import patch

with patch('deco.my_decorator.add', lambda function: function):
    from target import Target

class TestAdd:
    def test_add_decorator(self):
        sut = Target()
        actual = sut.execute_add()
        assert actual == 0

としてしまうと、

$ pytest test_global_patch_add.py 
...
test_global_patch_add.py:6: in <module>
    from target import Target
target.py:3: in <module>
    class Target:
target.py:15: in Target
    @add(2)
E   TypeError: 'int' object is not callable

 1 error in 0.22 seconds

エラーになります。

 
位置引数を持つデコレータのソースコードを見ると、引数なしのデコレータよりも、引数functionを持つ関数のネストが一段深くなっています。

def add(decorator_arg):
    def decorate(function):
        @wraps(function)
        def wrapped(*args, **kwargs):
            ...
            return result + decorator_arg
        return wrapped
    return decorate

これより、with patch('deco.my_decorator.add', lambda function: function):だと、本来はdecorate関数の引数functionを返すよう差し替えるべきなのに、add関数のdecorator_argを返すようになっていることが分かります。

その結果、callableなオブジェクトの受け取りを想定しているところで、intの2というcallableでないオブジェクトが返っているため、エラーとなっています。

 
そのため、テストコードで

with patch('deco.my_decorator.add', lambda decorator_arg: lambda function: function):
    from target import Target

と差し替えるlambdaを一段深くするように変更したところ、

$ pytest test_global_patch_add.py 
...
test_global_patch_add.py .

1 passed in 0.08 seconds

テストをパスしました。

 

ダミー処理をするデコレータに差し替える

今までは、何も処理しないデコレータへと差し替えました。

ただ、何らかの処理を行うデコレータに差し替えたいこともあります。

その場合は、

test_global_patch_replacement.py

from unittest.mock import patch
from functools import wraps

# 差し替え用のデコレータを用意
def dummy_decorator(function):
    @wraps(function)
    def wrapper(*args, **kwargs):
        result = function(*args, **kwargs)
        return result + 9
    return wrapper

# 用意したデコレータへ差し替える(引数無し版と引数あり版)
with patch('deco.my_decorator.countdown', lambda function: dummy_decorator(function)), \
        patch('deco.my_decorator.add', 
              lambda decorator_arg: lambda function: dummy_decorator(function)):
    from target import Target

class TestPatchArgs:
    def test_count_down_decorator(self):
        sut = Target()
        actual = sut.execute_count_down()
        assert actual == 9

    def test_add(self):
        sut = Target()
        actual = sut.execute_add()
        assert actual == 9

と、

  • 何らかの処理を行うデコレータを、デコレータの差し替え前に定義
    • 差し替え前に定義しないと「NameError: name ‘dummy_decorator’ is not defined」エラー
  • 何らかの処理を行うデコレータを返すようなlambdaにしてpatch

することで、

$ pytest test_global_patch_replacement.py 
...
test_global_patch_replacement.py ..

...
2 passed in 0.08 seconds

と、何らかの処理を行うデコレータへと差し替わりました。

 

デコレータの差し替え有無が混在する複数のテストファイルを同時に実行する場合

上記では、

  • デコレータを差し替えないテストコード
    • test_target.py
  • デコレータを差し替えるテストコード
    • test_global_patch_count_up.py

の2種類があります。

そこで、2つのテストファイルを同時に実行してみると、

$ pytest test_target.py test_global_patch_count_up.py 
...

test_target.py .....
test_global_patch_count_up.py F

...
self = <test_global_patch_count_up.TestCountUp object at 0x10b4cb710>

    def test_count_up_decorator(self):
        sut = Target()
        actual = sut.execute_count_up()
>       assert actual == 0
E       assert 1 == 0

test_global_patch_count_up.py:18: AssertionError
 1 failed, 5 passed in 0.13 seconds 

テストが失敗します。

テスト結果を見ると、デコレータを差し替えるテストコードにも関わらず、デコレータが差し替わっていないようです。

通常、複数のテストファイルを一括で実行してテストするため、このままでは使い物になりません。

 
そこでもう一度参考にしたブログを読むと、

I hope this silly example convinces you that decorators are applied when the class is created, which happens straight away when we load the module containing the class.

とあります。

ということは、複数のテストコードファイルでプロダクションコードがimportされた時に、デコレータの挙動が固定されてしまうのかなと考えられました。

 
ここで、Pythonでimportする場合、sys.modulesにモジュールに関するエントリが追加されることを思い出しました。

sys.modulesの挙動を詳しく見ると

新しいモジュールがインポートされたため、それらはsys.modules に追加される。これは、なぜ、同じモジュールを二回インポートする場合、それが非常に素早く行われるかを説明している: Python は、それをすでにロードし、sys.modules の中のモジュールとしてキャッシュしているため、二回目のインポートは単なる辞書への参照のみで済む。

6.4. sys.modules を使う - Dive into Python 5.4. (JAPANESE)

とのことです。

これより、importした時にデコレータの挙動が決まり、あとはそれが使い回されるのかなと考えました。

 
そこで、デコレータを差し替えるテストコードの場合のみ、事前にsys.modulesからエントリを削除してみます。

また、デコレータを差し替えない時のimportを優先させるよう、差し替えはテストメソッドの中で行うことにします。

test_method_patch_count_up.py

from unittest.mock import patch
import sys


class TestCountUpUsingWith:
    def test_count_up_decorator(self):
        # すでにimportされていることを考慮し、
        # そのtargetモジュールを無効化するためにsys.modulesより削除する
        # 未importの場合に例外KeyErrorとならないよう、第二引数にNoneを渡しておく
        sys.modules.pop('target', None)

        # 再度importしてパッチ
        with patch('deco.my_decorator.countup', lambda function: function):
            from target import Target

        # 検証
        sut = Target()
        actual = sut.execute_count_up()
        assert actual == 0

        # 使い終わったのでimportを削除
        sys.modules.pop('target')

 
テストしてみます。

$ pytest test_target.py test_method_patch_count_up.py 
...

test_target.py .....
test_method_patch_count_up.py .

6 passed in 0.12 seconds

想定通りの動きとなり、テストをパスしました。

 
なお、複数のテストケースでデコレータの差し替えを行いたい場合は、pytestのsetup/teardownを使い、

test_pytest_setup_patch_count_up.py

from unittest.mock import patch
import sys


class TestCountUpUsingWith:
    # 必要に応じて、クラスレベルやモジュールレベルにしても良い
    def setup_method(self):
        sys.modules.pop('target', None)

    def teardown_method(self):
        sys.modules.pop('target')

    def test_count_up_decorator(self):
        with patch('deco.my_decorator.countup', lambda function: function):
            from target import Target

        sut = Target()
        actual = sut.execute_count_up()
        assert actual == 0

としても

$ pytest test_target.py test_pytest_setup_patch_count_up.py 
...

test_target.py .....
test_pytest_setup_patch_count_up.py .

6 passed in 0.10 seconds

テストをパスします。

 

ソースコード

GitHubにあげました。e.g._mocking_decoratorディレクトリの中が今回のコードです。
thinkAmi-sandbox/python_mock-sample: Python : usage unittest.mock samples