Python testing using mock library for user input in a loop

Question

I am trying to use mock library for testing a piece of the code. In this code, the user raw input is accepted in a for loop as shown below. I have written the test case test_apple_record that can feed a single user input value for the tray number.

But, for every iteration within the for loop, it just takes this same value (5) as expected.

Question is: How to feed different values for each iteration? for example, specific values of 5, 6, and 7 for the tray numbers for i=0, 1 and 2 respectively.

class SomeClass(unittest.TestCase):    
    def apple_counter(self):
        apple_record = {}
        for i in range(3):
            apple_tray = input("enter tray number:")
            apple_record[apple_tray]  =  (i+1)*10
            print("i=%d, apple_record=%s"%(i, apple_record))

    def test_apple_record(self):
        with mock.patch('builtins.input', return_value='5'):
            self.apple_counter()

Answer 1

You can use the side_effect parameter with an iterable to provide return values:

with mock.patch('builtins.input', side_effect=[5, 6, 7]):
    self.apple_counter()

See the docs:

If side_effect is an iterable then each call to the mock will return the next value from the iterable.

Answer 2

Turns out I was reinventing the wheel. Use the side_effect parameter instead. Leaving a modified version of this here anyway, in case somebody wants to do something fancy.

To use a different function instead use the new_callable keyword argument of patch

with mock.patch('builtins.input', new_callable=lambda *x: random.randrange(0,10)):
    do_stuff()

however this means that values are entirely random, which is not desirable for a test because it means the test is not deterministic and can fail or not entirely by chance. Additionally you may want the input() function to emit generated non-random values, such as a sequence of commands or similar. For this, I think the easiest way would be a generator function.

def input_generator(): # generate squares as an example
    a = 0
    while True:
        yield a**2
        a += 1

g = input_generator()

with mock.patch('builtins.input', lambda *x: next(g)):
    do_stuff()