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

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开发者_JAVA技巧 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.

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)

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

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

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.