Python unittestで標準入出力をテストする方法

Pythonで標準入出力のテストを書く方法について説明します。ここでは、unittestとpytestの2つのフレームワークを使用します。

まず、テスト対象のコードを見てみましょう。以下のコードは、標準入力から受け取った2つの数字を足した答えを標準出力に表示する関数です。

def print_sum():
    x = input()
    y = input()
    print(int(x) + int(y))

次に、この関数をテストする方法を見てみましょう。

unittestでのテスト方法

unittestでの標準入出力のテストは以下のように行います。

import sys
import io
import unittest

def stub_stdin(testcase_inst, inputs):
    stdin = sys.stdin

    def cleanup():
        sys.stdin = stdin

    testcase_inst.addCleanup(cleanup)
    sys.stdin = io.StringIO(inputs)

def stub_stdouts(testcase_inst):
    stdout = sys.stdout

    def cleanup():
        sys.stdout = stdout

    testcase_inst.addCleanup(cleanup)
    sys.stdout = io.StringIO()

class MainTestCase(unittest.TestCase):
    def test_main(self):
        # 標準入力をモック
        stub_stdin(self, "1\\n4\\n")
        # 標準出力をモック
        stub_stdouts(self)
        print_sum()
        self.assertEqual(sys.stdout.getvalue(), "5\\n")

このコードでは、stub_stdinとstub_stdoutsで、sys.stdinとsys.stdoutをStringIOと入れ替えています。そして、addCleanupのメソッドを使い、テスト終了時には元のsys.stdinとsys.stdoutに戻るように設定しています。

Pytestでのテスト方法

Pytestではunittestに比べて大分シンプルに標準入力・標準出力のテストを書くことができます。

def test_main(capsys, monkeypatch):
    # 標準入力をモック
    monkeypatch.setattr('sys.stdin', io.StringIO("1\\n4\\n"))
    print_sum()
    # 標準出力のキャプチャを取得
    captured = capsys.readout()
    assert captured.out == "5\\n"

このコードでは、monkeypatchでsys.stdinをmockして、StringIOを返すようにしています。そしてcapsysで標準入力と標準出力のキャプチャを取得しています。

以上がPythonのunittestとpytestを使用した標準入出力のテスト方法です。これらの方法を使えば、標準入出力を含む関数のテストが容易になります。テストを書くことで、コードの品質を保つことができます。