Pythonic equivalent of ./foo.py < bar.png

Question

I've got a Python program that reads from sys.stdin, so I can call it with ./foo.py < bar.png. How do I test this code from within another Python module? That is, how do I set stdin to point to the contents of a file while running the test script? I don't want to do something like ./test.py < test.png. I don't think I can use fileinput, because the input is binary, and I only want to handle a single file. The file is opened using Image.open(sys.stdin) from PIL.

Answer 1

You should generalise your script so that it can be invoked from the test script, in addition to being used as a standalone program. Here's an example script that does this:

#! /usr/bin/python

import sys

def read_input_from(file):
    print file.read(),

if __name__ == "__main__":
    if len(sys.argv) > 1:
        # filename supplied, so read input from that
        filename = sys.argv[1]
        file = open(filename)
    else:
        # no filename supplied, so read from stdin
        file = sys.stdin
    read_input_from(file)

If this is called with a filename, the contents of that file will be displayed. Otherwise, input read from stdin will be displayed. (Being able to pass a filename on the command line might be a useful improvement for your foo.py script.)

In the test script you can now invoke the function in foo.py with a file, for example:

#! /usr/bin/python

import foo

file = open("testfile", "rb")
foo.read_input_from(file)

Answer 2

You can always monkey patch Your stdin. But it is quite ugly way. So better is to generalize Your script as Richard suggested.

import sys
import StringIO

mockin = StringIO.StringIO()
mockin.write("foo")
mockin.flush()
mockin.seek(0)

setattr(sys, 'stdin', mockin)

def read_stdin():
    f = sys.stdin
    result = f.read()
    f.close()
    return result

print read_stdin()

Also, do not forget to restore stdin when tearing down your test.

Answer 3

A comp.lang.python post showed the way: Substitute a StringIO() object for sys.stdout, and then get the output with getvalue():

def setUp(self):
    """Set stdin and stdout."""
    self.stdin_backup = sys.stdin

    self.stdout_backup = sys.stdout
    self.output_stream = StringIO()
    sys.stdout = self.output_stream

    self.output_file = None


def test_standard_file(self):
    sys.stdin = open(EXAMPLE_PATH)
    foo.main()
    self.assertNotEqual(
        self.output_stream.getvalue(),
        '')


def tearDown(self):
    """Restore stdin and stdout."""
    sys.stdin = self.stdin_backup
    sys.stdout = self.stdout_backup

Answer 4

• Your function or class should accept a stream instead of choosing which stream to use.
• Your main function will choose sys.stdin.
• Your test method will probably choose a StringIO instance or a test file.

The program:

# foo.py
import sys
from PIL import Image

def foo(stream):
  im = Image.open(stream)
  # ...

def main():
  foo(sys.stdin)

if __name__ == "__main__":
  main()

The test:

# test.py
import StringIO, unittest
import foo

class FooTest(unittest.TestCase):
  def test_foo(self):

    input_data = "...."
    input_stream = StringIO.StringIO(input_data)
    # can use a test file instead:
    # input_stream = open("test_file", "rb")

    result = foo.foo(input_stream)
    # asserts on result

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