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.