input()など標準入出力が使われた処理のテストにcontextlibの関数が便利かもしれない

はじめに

input() などユーザーの入力を受け取る系の処理が直書きされてしまっているコードはテストがしづらい。テストがしづらいということは毎回人の手でチェックしなければいけない。これではいけない。

入出力を伴う処理のテストについて

テストが出来るようにする方法は幾つかある。

1. 真面目に設計し直す
2. unittest.mockを使う
3. contextlibで提供されている標準入出力を奪う関数を使う

雑に対応できる汎用的な方法はunittest.mockを利用すること。ただしmockが必要になるのは設計の誤りなどと言われることもある。そうは言っても面倒な場合にはmockを使うこともあるかもしれない。

真面目に設計をし直す話は色んな所で見つけられると思う。今回は省略。

今回の話は3つ目についての話。mockでpathを細かく考えてpatchするよりもcontextlibの関数群を使う方が楽な場合もある。

標準出力を奪う

標準出力を奪う関数はデフォルトで用意されている contextlib.redirect_stdout。(同様に標準エラー出力を奪う関数 recirect_stderr も用意されている)

例えば print() が陽に使われた処理の出力部分のテストなどもStringIOを使って手軽に書くことが出来る。

import unittest
import contextlib


def foo():
    print("foo")


class InputTests(unittest.TestCase):
    def _calFUT(self):
        return foo()

    def test_it(self):
        from io import StringIO
        buf = StringIO()

        with contextlib.redirect_stdout(buf):
            self._calFUT()
        actual = buf.getvalue()
        self.assertEqual(actual, "foo\n")

標準入力を奪う

標準入力を奪う関数自体は提供されていない(つまるところイレギュラーな方法なのだけれど)。ただ、contextlib.redirect_stdout などの実装を見ることで意外と手軽に実装できる。

実際のところ contextlib.redirect_stdout などの実装は以下の様になっている。

class redirect_stdout(_RedirectStream):
    _stream = "stdout"

contextlib._RedirectStream はwithの前後で(__enter__と__exit__で)指定されたstreamの属性を入れ替えている。上の例では sys.stdout が入れ替えられる。これを使うことで標準入力を奪う redirect_stdin は手軽に実装できる。ただし _ から始まる名前の通りプライベートなオブジェクトなので将来に渡ってこの実装が動くかどうかは保証できない。

例えばパッケージ名を取得する関数の get_package_name() があるとする。

def get_package_name():
    package = input("input package name:")
    return {"package": package}

もちろん上で挙げた関数はテストのことが考えられていない良くない関数の定義の例ではあるけれど。これに対するテストは以下の様に書ける。おそらくそのままのかたちを残したままテストを書くのであればmockするよりも手軽。

import unittest
import contextlib


class redirect_stdin(contextlib._RedirectStream):
    _stream = "stdin"


class InputTests(unittest.TestCase):
    def _calFUT(self):
        return get_package_name()

    def test_it(self):
        from io import StringIO
        buf = StringIO()
        buf.write("hello\n")
        buf.seek(0)

        with redirect_stdin(buf):
            actual = self._calFUT()

        expected = {"package": "hello"}
        self.assertEqual(actual, expected)

おまけ

contextlib.redirect_stdout を使った処理のサンプルとしてこういうような処理も書けるる。通常の出力結果にインデントを付加した環境を用意したい場合などに便利かもしれない。

print("a")
with indent(2):
    print("b")
    with indent(2):
        print("c")
    print("d")
print("e")

これは以下のような出力になる。

a
  b
    c
  d
e

実装は以下のようなもの。

import sys
import contextlib
from io import StringIO


@contextlib.contextmanager
def indent(n):
    buf = StringIO()
    with contextlib.redirect_stdout(buf):
        yield buf
    buf.seek(0)

    prefix = " " * n
    write = sys.stdout.write
    for line in buf:
        write(prefix)
        write(line)
    sys.stdout.flush()