Pythonのunittestで標準出力のテストを作る

最近やたらunittestを書いているのですが、標準出力部分のテストをしなければいけなくなり、やり方を確認したのでメモとして残します。

どういうことか

こういう標準出力に内容を書き出す関数についてunittestを作成したかったという話です。

def hoge():
    print("HELLO")

ポイント

print等で書き出される標準出力はsys.stdoutだということ、そしてsys.stdoutは再代入可能という点がポイントです。各テストケース前にsys.stdoutをio.StringIO()等のファイルライクなオブジェクトで置き換えてしまえばテストができるようになります。ただ、そのまま上書きしてしまうと支障が出るのでオリジナルのstdoutも退避しておきます。

    def setUp(self):
        """
        :return:
        """
        # sys.stdoutを退避し、StringIO()で上書きする
        self.org_stdout, sys.stdout = sys.stdout, StringIO()

    def tearDown(self):
        """
        :return:
        """
        # 退避していたsys.stdoutを戻す
        sys.stdout = self.org_stdout

余談ですがa, b = b, aみたいなスワップが1行でかけるのはいいですね。

コード

実際のテストコードはこんな感じになります。

# coding: utf-8
import sys
import unittest
from io import StringIO


# テスト対象のコード
def hoge():
    print("HELLO")

class TestHoge(unittest.TestCase):

    def setUp(self):
        """
        :return:
        """
        # sys.stdoutを退避し、StringIO()で上書きする
        self.org_stdout, sys.stdout = sys.stdout, StringIO()

    def tearDown(self):
        """
        :return:
        """
        # 退避していたsys.stdoutを戻す
        sys.stdout = self.org_stdout
    
    def test_01(self):
        hoge()
        # HELLOが書き出されているか.printなので改行が付与されている
        self.assertEqual(sys.stdout.getvalue(), "HELLO\n")

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

実行するとこんな感じです。

DOS> python test.py
test_01 (__main__.TestHoge) ... ok

----------------------------------------------------------------------
Ran 1 test in 0.001s

OK

ちゃんとテストできていますね。