Python - Mock input() method

test_mock_input.py

from unittest.mock import patch

@patch('builtins.input', lambda _: 'y') # <- value to be returned by input method
def test_should_be_the_correct_value(self):

    # replace with your class or method to be tested:
    value = input('enter with a value')

    assert value == 'y'

test_mock_input_2.py

# Other way to make the same mock:

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

def test_should_be_the_correct_value(self):
    with MockInputFunction('y'): # <- value to be returned by input method
        
        # replace with your class or method to be tested:
        value = input('enter with a value')
        
        assert value == 'y'