pythonのdecoratorをmock関数で置き換える方法

結論

decoratorをモックや別の関数で置き換えるときはもう一度importしなおそう！

これで理解できた方はこの後の説明は不要です笑
よくわからんって方は解説をご覧ください

解説

テストするときに特定の処理のみをモックアップしたり置き換えたりしたいってことよくあると思います！

Easy Case🐕

例えばこういうケース

def gopher():
	print("Goに入ってはGoに従え")
	inner_function()
	return 0

def inner_function():
	print("Inner")

このように内部で使われる関数をモックアップするだけなら@patchデコレーターを使うだけで、対象の関数をモックアップできます！簡単ですね

from unittest import mock
from somemodule import gopher

mock.patch("inner_function")
def test_gopher()
	assert gopher() == 0

Tricky Case🐍

ではこのケースはいかがでしょう？DBの接続をデコレーターでラッピングすることでコードの再利用性を高めたものです。

# db.py
def use_db(func):
	def inner():
		print("Use Real DB")
		# ここに本番用DBの接続コードなどがある
		# MySQLやPostgreSQLなど
		return func()
	return inner

# somemodule.py
from db import use_db
	
@use_db
def gopher():
	print("Goに入ってはGoに従え")
	return 0

テストコードでは本番用のDBでテストするのはやばいので、今回はテスト用のDB（sqlite3など）に切り替えれるようにしましょう。_test_db()という置き換え用の関数を用意してデコレーターを置き換えてみます。
pythonの関数は基本的には下記の様に再代入することで置換可能なので、まずはそれを試してみます。

間違った例🙅‍♂️

# test_somemodule.py
from unittest import mock
from somemodule import gopher

def _test_db():
	def inner():
		print("Use Test DB")
		# ここにテストDBの接続コードなどがある
		return func()
	return inner	

@mock.patch("db.use_db", _test_db)
def test_gopher():
	use_db = _test_db
	assert gopher() == 0

出力

Goに入ってはGoに従え
Use Real DB

おや？なぜか元のuse_dbが呼ばれてますね...
本来はここで"Use Test DB"というが表示されてほしいはずです。ではどうやればいいでしょうか？

まず原因としては、gopherモジュールをimportした瞬間に@use_dbデコレーターは元のuse_db関数として読み込まれます。 なので、モックアップである_test_dbをgopherが読み込まれる前にパッチを適用する必要があります。 patchの際はstart()を忘れない様にしてください！

＊今回のようにuse_dbはからなずsomemodule.pyとは別ファイル(db.py)においてください。そうしないと上書きできなくなります。

# test_somemodule.py
from unittest import mock

def _test_db(func):
    def inner():
        print("Use Test DB")
        # ここにテストDBの接続コードなどがある
        return func()

    return inner

mock.patch("db.use_db", _test_db).start()
from somemodule import gopher

def test_gopher():
    assert gopher() == 0

出力例

Use Test DB
Goに入ってはGoに従え

いかがでしたか？これでテスト用のモックアップも怖くないはずです！今回のサンプルコードはこちらにあげましたので実際に動かしてみたい方はご利用ください！

また、デコレーターをpatchで置き換える他の方法は参考リンクの記事に書いてあるので、興味ある方は合わせてご覧ください。
それでは良いpythonライフを！🐍🐍🐍

参考

Can I patch a Python decorator before it wraps a function?
Python Mock Gotchas|Alex Marandon