Showing posts with label raw_input. Show all posts
Showing posts with label raw_input. Show all posts

Thursday, July 11, 2013

Monkey patch (mock) the built-in raw_input function

For testing without mock (thirdparty or 3.x), I wrote a mock context manager function for raw_input. If you do not provide enough response in the mock_raw_input argument list, the raw_input is back to normal.
import __builtin__
import sys
from contextlib import contextmanager

@contextmanager
def mock_raw_input(*input_list):
    """Monkey patch the raw_input()"""
    _raw_input = __builtin__.raw_input
    def _stub(prompt='', input_iter=iter(input_list), orig=_raw_input):
        try:
            _input = next(input_iter)
            sys.stdout.write(prompt)
            return _input
        except StopIteration:
            __builtin__.raw_input = _raw_input
            return raw_input(prompt)
    __builtin__.raw_input = _stub
    try:
        yield
    finally:
        __builtin__.raw_input = _raw_input

def main():
    with mock_raw_input('abc','cde'):
        print raw_input('abc?\n')
        print raw_input()
        print raw_input('no more expected value? ')
    raw_input('raw_input should be back to default here, right? ')

if __name__ == "__main__":
    main()