Simple Python unittest for mocking/stubbing sys.stdin and friends using native libs.

test_stdio.py

# -*- coding: utf-8 -*-
# To use, at the directory you saved this:
# python -m unittest test_stdio

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 = StringIO(inputs)


def stub_stdouts(testcase_inst):
    stderr = sys.stderr
    stdout = sys.stdout

    def cleanup():
        sys.stderr = stderr
        sys.stdout = stdout

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


class StdioTestCase(unittest.TestCase):

    def test_example(self):
        stub_stdin(self, '42')
        stub_stdouts(self)
        example()
        self.assertEqual(sys.stdout.getvalue(), '42\n')

    def helper(self, data, answer, runner):
        stub_stdin(self, data)
        stub_stdouts(self)
        runner()
        self.assertEqual(sys.stdout.getvalue(), answer)
        # self.doCleanups()  # optional, see comments below

    def test_various_inputs(self):
        data_and_answers = [
            ('hello', 'HELLOhello'),
            ('goodbye', 'GOODBYEgoodbye'),
        ]

        runScript = upperlower  # the function I want to test

        for data, answer in data_and_answers:
            self.helper(data, answer, runScript)


class StringIO(io.StringIO):
    """
    A "safely" wrapped version
    """

    def __init__(self, value=''):
        value = value.encode('utf8', 'backslashreplace').decode('utf8')
        io.StringIO.__init__(self, value)

    def write(self, msg):
        io.StringIO.write(self, msg.encode(
            'utf8', 'backslashreplace').decode('utf8'))


def example():
    number = raw_input()
    print number.upper()


def upperlower():
    raw = raw_input()
    print (raw.upper() + raw),


# Calling doCleanup is optional simply because TestCase will call it by
# default.  It's there to clean up instantly if it is desired and won't
# invoke other cleanup code prematurely.

Hi @metatoaster! Would you be willing to release this code under the MIT license?

@lvreynoso sure, I don't mind, given that this trivial piece of code was originally created to demonstrate to some of my colleagues/friends on how redirection of stdio might be done for testing in Python 2.7, while being compatible with Python 3. However, I honestly don't recommend this solution any more, as the implementation has a tight coupling with the unittest framework (the function argument depends directly on an instance of unittest.TestCase), and that Python 2 has been fully deprecated.

A better solution going forward would be to implement a context manager that manage the stdio redirection, where an example of how this might be done may be found at issue15805 on the CPython issue tracker. Alternatively, unittest.mock.patch may be leveraged to achieve what the above example code does, given that the linked documentation provide an example to achieve the redirection of sys.stdout. If both sys.stdin and sys.stdout are needed, the following example shows how it might be done in Python 3:

from io import StringIO
from unittest.mock import patch

with patch('sys.stdin', StringIO('Darcy\n')) as stdin, \
         patch('sys.stdout', new_callable=StringIO) as stdout:
    name = input('What is your name? ')
    print('Hello %s' % name)

    assert stdout.getvalue() == 'What is your name? Hello Darcy\n'
    assert stdin.read() == ''  # all input consumed

Yes, please feel free to use this additional trivial code example should this Python 3.3+ solution is more preferable to your needs.

Thank you so much! I appreciate your thoughtful comments.

You're very welcome! (also apologies for getting your name wrong initially when doing the @ reply)