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:
raw_input
(an input side effect) needs to be mocked.
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 yield
ed 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.