python mocking raw input in unittests

You can't patch input but you can wrap it to use mock.patch(). Here is a solution:

from unittest.mock import patch
from unittest import TestCase


def get_input(text):
    return input(text)


def answer():
    ans = get_input('enter yes or no')
    if ans == 'yes':
        return 'you entered yes'
    if ans == 'no':
        return 'you entered no'


class Test(TestCase):

    # get_input will return 'yes' during this test
    @patch('yourmodule.get_input', return_value='yes')
    def test_answer_yes(self, input):
        self.assertEqual(answer(), 'you entered yes')

    @patch('yourmodule.get_input', return_value='no')
    def test_answer_no(self, input):
        self.assertEqual(answer(), 'you entered no')

Keep in mind that this snippet will only work in Python versions 3.3+

Okay, first off, I feel it's necessary to point out that in the original code in question, there are actually two things that need to be tackled:

1. raw_input (an input side effect) needs to be mocked.
2. print (an output side effect) needs to be checked.

In an ideal function for unit testing, there would be no side effects. A function would simply be tested by handing in arguments and its output would be checked. But often we want to test functions which aren't ideal, IE, in functions like yours.

So what are we to do? Well, in Python 3.3, both of the issues I listed above became trivial because the unittest module gained the ability to mock and check for side effects. But, as of the start of 2014, only 30% of Python programmers had moved on to 3.x, so for the sake of the other 70% of Python programmers still using 2.x, I'll outline an answer. At the current rate, 3.x won't overtake 2.x until ~2019, and 2.x won't vanish until ~2027. So I figure this answer will be useful for several years to come.

I want to address the issues listed above one at a time, so I'm going to initially change your function from using print as its output to using return. No surprises, here's that code:

def answerReturn():
    ans = raw_input('enter yes or no')
    if ans == 'yes':
        return 'you entered yes'
    if ans == 'no':
        return 'you entered no'

So all we need to do is mock raw_input. Easy enough - Omid Raha's answer to this very question shows us how to do that by swizzling out the __builtins__.raw_input implementation with our mock implementation. Except his answer wasn't properly organized into a TestCase and functions, so I'll demonstrate that.

import unittest    

class TestAnswerReturn(unittest.TestCase):
    def testYes(self):
        original_raw_input = __builtins__.raw_input
        __builtins__.raw_input = lambda _: 'yes'
        self.assertEqual(answerReturn(), 'you entered yes')
        __builtins__.raw_input = original_raw_input

    def testNo(self):
        original_raw_input = __builtins__.raw_input
        __builtins__.raw_input = lambda _: 'no'
        self.assertEqual(answerReturn(), 'you entered no')
        __builtins__.raw_input = original_raw_input

Small note just on Python naming conventions - variables which are required by the parser but not used are typically named _, as in the case of the lambda's unused variable (which is normally the prompt shown to the user in the case of the raw_input, incase you're wondering why it's required at all in this case).

Anyways, this is messy and redundant. So I'm going to do away with the repetition by adding in a contextmanager, which will allow for simple with statements.

from contextlib import contextmanager

@contextmanager
def mockRawInput(mock):
    original_raw_input = __builtins__.raw_input
    __builtins__.raw_input = lambda _: mock
    yield
    __builtins__.raw_input = original_raw_input

class TestAnswerReturn(unittest.TestCase):
    def testYes(self):
        with mockRawInput('yes'):
            self.assertEqual(answerReturn(), 'you entered yes')

    def testNo(self):
        with mockRawInput('no'):
            self.assertEqual(answerReturn(), 'you entered no')

I think that nicely answers the first part of this. On to the second part - checking print. I found this much trickier - I'd love to hear if anyone has a better answer.

Anyways, the print statement can't be overridden, but if you use print() functions instead (as you should) and from __future__ import print_function you can use the following:

class PromiseString(str):
    def set(self, newString):
        self.innerString = newString

    def __eq__(self, other):
        return self.innerString == other

@contextmanager
def getPrint():
    promise = PromiseString()
    original_print = __builtin__.print
    __builtin__.print = lambda message: promise.set(message)
    yield promise
    __builtin__.print = original_print

class TestAnswer(unittest.TestCase):
    def testYes(self):
        with mockRawInput('yes'), getPrint() as response:
            answer()
            self.assertEqual(response, 'you entered yes')

    def testNo(self):
        with mockRawInput('no'), getPrint() as response:
            answer()
            self.assertEqual(response, 'you entered no')

The tricky bit here is that you need to yield a response before the with block is entered. But you can't know what that response will be until the print() inside the with block is called. This would be fine if strings were mutable, but they aren't. So instead a small promise or proxy class was made - PromiseString. It only does two things - allow a string (or anything, really) to be set and let us know if it's equal to a different string. A PromiseString is yielded and then set to the value that would normally be print within the with block.

Hopefully you appreciate all this trickery I've written up since it took me around 90 minutes to put together this evening. I tested all of this code and verified it all worked with Python 2.7.

I am using Python 3.4 and had to adapt answers above. My solution factors out common code into the custom runTest method and shows you how to patch both input() and print(). Here's code that works as advertised:

import unittest
from io import StringIO
from unittest.mock import patch


def answer():
    ans = input('enter yes or no')
    if ans == 'yes':
        print('you entered yes')
    if ans == 'no':
        print('you entered no')


class MyTestCase(unittest.TestCase):
    def runTest(self, given_answer, expected_out):
        with patch('builtins.input', return_value=given_answer), patch('sys.stdout', new=StringIO()) as fake_out:
            answer()
            self.assertEqual(fake_out.getvalue().strip(), expected_out)

    def testNo(self):
        self.runTest('no', 'you entered no')

    def testYes(self):
        self.runTest('yes', 'you entered yes')

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