[PYTHON] Contextlib functions may be useful for testing processes that use standard I / O such as input ()

Introduction

It is difficult to test the code where the processing of the system that receives the user's input such as ʻinput ()` is written directly. Every time you check that the test is difficult, you have to check it manually. This should not be.

About testing processing involving input and output

There are several ways to enable testing.

  1. Seriously redesign
  2. Use unittest.mock
  3. Use the standard I / O deprivation function provided by contextlib

Use unittest.mock as a general-purpose method that can handle miscellaneous things. However, it is sometimes said that the need for mock is a design error. That said, you might use a mock if it's a hassle.

I think you can find stories about seriously redesigning in various places. Omitted this time.

This story is about the third one. Sometimes it's easier to use contextlib's functions than to patch the path in detail with a mock.

Steal standard output

The function that steals standard output is provided by default contextlib.redirect_stdout. (Similarly, there is also a function recirect_stderr that steals standard error)

For example, you can easily write a test of the output part of the process in which print () is explicitly used using 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")

Steal standard input

The function itself that steals standard input is not provided (although it is an irregular method after all). However, it is surprisingly easy to implement by looking at the implementation such as contextlib.redirect_stdout.

Actually, the implementation such as contextlib.redirect_stdout is as follows.

class redirect_stdout(_RedirectStream):
    _stream = "stdout"

contextlib._RedirectStream swaps the attributes of the specified stream before and after with (with __enter__ and __exit__). In the above example, sys.stdout is replaced. You can easily implement redirect_stdin, which steals standard input by using this. However, as the name starts with _, it is a private object, so we cannot guarantee that this implementation will work in the future.

For example, suppose you have a function get_package_name () that gets the package name.

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

Of course, the functions listed above are examples of bad function definitions that aren't considered for testing. The test for this can be written as follows. Probably easier than mocking if you want to write a test while leaving it as it is.

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)

bonus

You can also write such a process as a sample of the process using contextlib.redirect_stdout. It may be convenient when you want to prepare an environment with indentation added to the normal output result.

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

This will produce the following output.

a
  b
    c
  d
e

The implementation is as follows.

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()

Recommended Posts

Contextlib functions may be useful for testing processes that use standard I / O such as input ()
A collection of resources that may be useful for creating and expanding dotfiles