How to test a function with input call?

Viewed 44330

I have a console program written in Python. It asks the user questions using the command:

some_input = input('Answer the question:', ...)

How would I test a function containing a call to input using pytest? I wouldn't want to force a tester to input text many many times only to finish one test run.

8 Answers

This can be done with mock.patch and with blocks in python3.

import pytest
import mock
import builtins

"""
The function to test (would usually be loaded
from a module outside this file).
"""
def user_prompt():
    ans = input('Enter a number: ')
    try:
        float(ans)
    except:
        import sys
        sys.exit('NaN')
    return 'Your number is {}'.format(ans)

"""
This test will mock input of '19'
"""    
def test_user_prompt_ok():
    with mock.patch.object(builtins, 'input', lambda _: '19'):
        assert user_prompt() == 'Your number is 19'

The line to note is mock.patch.object(builtins, 'input', lambda _: '19'):, which overrides the input with the lambda function. Our lambda function takes in a throw-away variable _ because input takes in an argument.

Here's how you could test the fail case, where user_input calls sys.exit. The trick here is to get pytest to look for that exception with pytest.raises(SystemExit).

"""
This test will mock input of 'nineteen'
"""    
def test_user_prompt_exit():
    with mock.patch.object(builtins, 'input', lambda _: 'nineteen'):
        with pytest.raises(SystemExit):
            user_prompt()

You should be able to get this test running by copy and pasting the above code into a file tests/test_.py and running pytest from the parent dir.

Since I need the input() call to pause and check my hardware status LEDs, I had to deal with the situation without mocking. I used the -s flag.

python -m pytest -s test_LEDs.py

The -s flag essentially means: shortcut for --capture=no.

A different alternative that does not require using a lambda function and provides more control during the tests is to use the mock decorator from the standard unittest module.

It also has the additional advantage of patching just where the object (i.e. input) is looked up, which is the recommended strategy.

# path/to/test/module.py
def my_func():
    some_input = input('Answer the question:')
    return some_input
# tests/my_tests.py

from unittest import mock

from path.to.test.module import my_func

@mock.patch("path.to.test.module.input")
def test_something_that_involves_user_input(mock_input):
    mock_input.return_value = "This is my answer!"
    assert my_func() == "This is my answer!"
    mock_input.assert_called_once()  # Optionally check one and only one call   

You can also use environment variables in your test code. For example if you want to give path as argument you can read env variable and set default value if it's missing.

import os
...
input = os.getenv('INPUT', default='inputDefault/')

Then start with default argument

pytest ./mytest.py

or with custom argument

INPUT=newInput/ pytest ./mytest.py
Related