Search code examples
pythonunit-testingpython-3.xlambdapatch

Mock standard input - multi line in python 3


I am new to python and have been using python 3 for my learning. I am using python's unit test framework to test my code.

Problem :-

The function that I need to unit test takes inputs in the following manner:-

def compare():
   a, b, c = input().strip().split(' ')
   d, e, f = input().strip().split(' ')
   # other code here

I am using the following test case to mock the input :-

class TestCompare(unittest.TestCase):

   @patch("builtins.input", lambda: "1 2 3")
   @patch("builtins.input", lambda: "4 5 6")
   def test_compare(self):
      self.assertEqual(compare(), "1 1")

The problem I am facing is that when the test case is run the variable triplets a,b,c and d,e,f have the same values - 1,2,3

I have been trying to find a way to inject the second set of inputs to run my test but in vain.

Any help regarding the above is greatly appreciated.

Solution environment :- Python 3


Solution

  • The patch decorator will ensure the patched function always return that value, and if subsequent calls must be different, your mock object must have a way to simulate that. This ends up being much more complicated.

    What you can do however is go one step lower and patch the underlying layer, which is the standard input/output layer. One common strategy that other test frameworks have done is to deal with the sys.stdin and sys.stdout objects directly. Consider this:

    import unittest
    from unittest.mock import patch
    
    from io import StringIO
    
    def compare():
        a, b, c = input().strip().split(' ')
        d, e, f = input().strip().split(' ')
    
        return '%s %s' % (a, d)
    
    class TestCompareSysStdin(unittest.TestCase):
    
        @patch("sys.stdin", StringIO("1 2 3\n4 5 6"))
        def test_compare(self):
            self.assertEqual(compare(), "1 4")
    

    Execution

    $ python -m unittest foo
    .
    ----------------------------------------------------------------------
    Ran 1 test in 0.000s
    
    OK
    

    Naturally, this works at a lower level, and so the option to have an iterator that returns different values on subsequent calls may be more suitable.