python mocking raw input in unittests

Suppose I have this python code:

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

How do I write a unittest for this ? I know i have to use 'Mock' but I don't understand how. Can anyone make some simple example ?

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

Just ran across the same problem, but I just mocked out __builtin__.raw_input.

Only tested on Python 2. pip install mock if you don't already have the package installed.

from mock import patch
from unittest import TestCase

class TestAnswer(TestCase):
    def test_yes(self):
        with patch('__builtin__.raw_input', return_value='yes') as _raw_input:
            self.assertEqual(answer(), 'you entered yes')
            _raw_input.assert_called_once_with('enter yes or no')

    def test_no(self):
        with patch('__builtin__.raw_input', return_value='no') as _raw_input:
            self.assertEqual(answer(), 'you entered no')
            _raw_input.assert_called_once_with('enter yes or no')

Alternatively, using the library genty, you can simplify the two tests:

from genty import genty, genty_dataset
from mock import patch
from unittest import TestCase

@genty
class TestAnswer(TestCase):
    @genty_dataset(
        ('yes', 'you entered yes'),
        ('no', 'you entered no'),
    )
    def test_answer(self, expected_input, expected_answer):
        with patch('__builtin__.raw_input', return_value=expected_input) as _raw_input:
            self.assertEqual(answer(), expected_answer)
            _raw_input.assert_called_once_with('enter yes or no')

Here's what I do in Python 3:

class MockInputFunction:
    def __init__(self, return_value=None):
        self.return_value = return_value
        self._orig_input_fn = __builtins__['input']

    def _mock_input_fn(self, prompt):
        print(prompt + str(self.return_value))
        return self.return_value

    def __enter__(self):
        __builtins__['input'] = self._mock_input_fn

    def __exit__(self, type, value, traceback):
        __builtins__['input'] = self._orig_input_fn

which can then be used in any context. For instance, pytest uses ordinary assert statements.

def func():
    """ function to test """
    x = input("What is x? ")
    return int(x)

# to test, you could simply do:
with MockInputFunction(return_value=13):
    assert func() == 13

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


def test_answer_yes():
    assert(answer() == 'you entered yes')

def test_answer_no():
    assert(answer() == 'you entered no')

origin_raw_input = __builtins__.raw_input

__builtins__.raw_input = lambda x: "yes"
test_answer_yes()

__builtins__.raw_input = lambda x: "no"
test_answer_no()

__builtins__.raw_input = origin_raw_input

This question was the best source of solving my mocking of builtins.input need in 2022. I've expanded Kris's answer for more functionality that better fit my needs. Specifically, I need side effects, if only in the form of a (optional) list:

from unittest.mock import call

class MockInputFunction:
    def __init__(self, return_value=None, side_effect=None):
        self.return_value = return_value
        self.side_effect = side_effect
        self.mock_calls = []
        self._orig_input_fn = __builtins__['input']

    def _mock_input_fn(self, prompt=None):
        return_value = self.return_value\
            if self.side_effect is None\
            else self.side_effect[len(self.mock_calls)]
        self.mock_calls.append(call(prompt))
        if prompt:
            print(prompt + str(return_value))
        else:
            print(str(return_value))
        return return_value

    def __enter__(self):
        __builtins__['input'] = self._mock_input_fn

    def __exit__(self, type, value, traceback):
        __builtins__['input'] = self._orig_input_fn

In action:

def some_func():
    input()
    binstr = input()
    changes= 1
    # Clever stuff
    return changes


def test_case1():
    with MockInputFunction(side_effect=["","1101110"]):
        changes = some_func()
    print(changes)
    assert changes == 1