unittest の addModuleCleanup/doModuleCleanups 関数の紹介

pythonのunittest,pytest周りをお勉強中に公式日本語訳のunittestのドキュメントをみてたら、
python3.8から追加されたaddModuleCleanup、doModuleCleanups関数があったので紹介しときます。
unittestって何？って方は、この記事を読むより、別の記事でunittestの概要を掴んだほうがいいかもしれません。

TL;DR

setUpModule関数でエラー落ちしたときに呼ぶ関数指定できますよ関数

unittest復習

例えば以下のコードを実行してみます。
足し算と引き算を確認するテストです。
どのタイミングで関数が実行されているかを確認するため、
print(sys._getframe().f_code.co_name)にて関数名を出力しています。

import sys
import unittest

# TODO(hatsumi) ～後でコード追加します～

def createConnection():
    print(sys._getframe().f_code.co_name)
#    raise Exception

def closeConnection():
    print(sys._getframe().f_code.co_name)

def setUpModule():
    createConnection()

def tearDownModule():
    closeConnection()

class TestOperator(unittest.TestCase):

    @classmethod
    def setUpClass(self):
        print(sys._getframe().f_code.co_name)

    @classmethod
    def tearDownClass(self):
        print(sys._getframe().f_code.co_name)

    def setUp(self):
        print(sys._getframe().f_code.co_name)

    def tearDown(self):
        print(sys._getframe().f_code.co_name)

    def test_add(self):
        print(sys._getframe().f_code.co_name)
        self.assertEqual(2 + 3, 5)

    def test_sub(self):
        print(sys._getframe().f_code.co_name)
        assert 3 - 2 == 1


if __name__ == "__main__":
    unittest.main()

実行結果として、以下が出力されるはずです。

createConnection
setUpClass
setUp
test_add
tearDown
.setUp
test_sub
tearDown
.tearDownClass
closeConnection

module -> class -> methodの順にsetup,teardownが呼ばれている事がわかります。

本題

setUpModule関数で、exceptionが発生してしまった場合、tearDownModule関数は呼び出されません。
その代わりにaddModuleCleanup関数で追加された関数が、doModuleCleanups関数によって実行されます。

上記コードでは、setUpModule関数にてcreateConnection、tearDownModule関数でcloseConnectionと出力しています。
setUpModule関数のコメントを外してraise Exceptionが実行されると、tearDownModule関数は呼ばれません。
データベースコネクションだったり、テンポラリなものの処理を、tearDownModule関数が呼ばれる前提で管理しているならばいいかもしれません。

実装

1つの関数の登録

関数を実行するためには、事前に関数を登録しておく必要があります。
そのためにはaddModuleCleanup関数を使用します。
上記コードのTODOコメント部分に以下のコードを追記します。
print_number関数を登録しておきます。

def print_number():
    print(sys._getframe().f_code.co_name)

unittest.addModuleCleanup(print_number)

実行結果としてExceptionが発生すると、print_number関数が呼び出されています。
createConnectionも文字列は出ちゃってますね。。

print_number
createConnection
E
======================================================================
ERROR: setUpModule (__main__)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "test_module.py", line 25, in setUpModule
    createConnection()
  File "test_module.py", line 19, in createConnection
    raise Exception
Exception

----------------------------------------------------------------------
Ran 0 tests in 0.001s

複数の関数の登録

複数の関数を登録しておくことも可能です。
ただその場合は、LIFOで実行されます。
上記コードのTODOコメント部分に以下のコードを追記します。

unittest.addModuleCleanup((lambda : print('a')))
unittest.addModuleCleanup((lambda : print('b')))
unittest.addModuleCleanup((lambda : print('c')))

実行結果

c
b
a
createConnection
E

a->b->cの順で登録して、c->b->aの順で出力されています。

登録した関数の実行

登録した関数はdoModuleCleanups関数によって実行されます。
鋭い方は気づかれてるかもしれませんが、LIFOで実行する際には、listからpopしています。
なのでdoModuleCleanups関数を先に実行しておけば、実行する関数リストは殻なので何も実行されません。

上記コードのTODOコメント部分に以下のコードを追記します。

unittest.addModuleCleanup((lambda : print('a')))
unittest.addModuleCleanup((lambda : print('b')))
unittest.addModuleCleanup((lambda : print('c')))

unittest.case.doModuleCleanups()
print()

出力結果

c
b
a

createConnection
E

　print()によって1行の空きがあったあとに、Exceptionが発生しています。
再度登録されている関数が実行されることはありません。
またdoModuleCleanups関数なのですが、公式ドキュメントではcaseが省略されていますので、実際にはunittest.case.doModuleCleanups()と呼ぶ必要があります。（誤記？）

まとめと余談

個別で記事がなかったはず(2020/5/4時点)なので作成してみました。
pytestではこの機能はない感じですかね・・
お役に立てば幸甚です。

心優しい方は以下の質問にも答えて下さい←
・
・
・

複数の関数を登録する際に最初は以下のように書いてたんですが、
idが同じになってしまったので、記事のようにabc毎にlambdaを出力しました。
どうすればインスタンスを個別に生成できますか・・？（ワンライナーじゃむりぽですか。。）

#これだと全部4が出力される・・
for i in range(5):
    unittest.addModuleCleanup((lambda : print(i)))

URL

• 詳しい説明は公式へどうぞ　→　https://docs.python.org/ja/3/library/unittest.html