Automate Stdin For Python Tests

NOTE: The information in this post applies to Python 2.x only. See this post for Python 3.x support.

While working on Qprompt (a Python CLI based user input libary), I wanted to write some tests to check for regressions. The following lambda function allows stdin to be programmatically set:

import StringIO, sys
sys.stdin = StringIO.StringIO()
setinput = lambda x: [sys.stdin.truncate(0), sys.stdin.write(x), sys.stdin.seek(0)]

The following is an example test scenario that checks the ask_int() function:

setinput("42")
result = ask_int()
print result        # 42
print type(result)  # <type 'int'>