April 25, 2022

Unit-testing the command line input in Python

As essential a practice as unit testing is in everyday development, it can be puzzling in the presence of prompt inputs. And the matter gets even hairier for passwords. Let's see how to effectively handle every such case easily.

For the purpose of this demonstration we'll use the renown click module.

Let's say we have a function defining prompts for both a username and a password.

import click

def collect_user_data():
    """Prompt for and collect user data."""
    username = click.prompt("Enter your desired username")
    password = click.prompt("Enter your desired password", hide_input=True)
    return username, password

We write a unit test and draft our input mock idea.

from unittest import TestCase, mock

@contextmanager
def input(*cmds):
    """Replace input."""
    cmds = "\n".join(cmds)
    with mock.patch("sys.stdin", StringIO(f"{cmds}\n"))):
        yield

class TestInput(TestCase):
    """Test the prompt input."""

    def test_collect_user_data(self):
        """Test returning the user data collected from prompt."""
        with input("dummy.username", "p4sSw0rD"):
            self.assertEqual(collect_user_data(), ("dummy.username", "p4sSw0rD"))

Looks fine. Mocking the standard input's return value requires commands to be passed together as a string stream, with each command separated by a new line character, representing the "Enter" key. Let's run the test and see what happens.

$ python3 -m unittest
Enter your desired username: Enter your desired password:

This doesn't bode well: the test routine hangs at the password input. There's clearly something we're missing in our mocking logic.

Beyond its documentation, digging deeper in its code, we realize the click module handles password prompts differently from regular visible inputs.

Such hidden inputs are aptly implemented using the Python's built-in getpass module, which omits echoing the typed-in characters. It also acts as a convenient abstraction layer across both the Unix and Windows operating systems - hence its claimed portability.

We should therefore introduce a new logic branch dealing with the aforementioned module, to tell visible inputs from hidden ones in our mock.

We can get as creative as we want here: we choose to avoid unnecessarily messing with the internals of getpass and return the hidden values as callable side_effects and pass a different data type for hidden inputs with some kind of placeholder key.

Let's tweak our mock function and test accordingly:

from unittest import TestCase, mock

@contextmanager
def input(*cmds):
    """Replace input."""
    visible_cmds = "\n".join([c for c in cmds if isinstance(c, str)])
    hidden_cmds = [c.get("hidden") for c in cmds if isinstance(c, dict)]
    with mock.patch("sys.stdin", StringIO(f"{visible_cmds}\n")), 
        mock.patch("getpass.getpass", side_effect=hidden_cmds):
            yield

class TestInput(TestCase):
    """Test the prompt input."""

    def test_collect_user_data(self):
        """Test returning the user data collected from prompt."""
        with input("dummy.username", {"hidden" : "p4sSw0rD"}):
            self.assertEqual(collect_user_data(), ("dummy.username", "p4sSw0rD"))

Running our test again now rewards us with the sought-for success message.

$ python3 -m unittest 
Enter your desired username: Enter your desired password:.
----------------------------------------------------------------------
Ran 1 test in 0.003s

OK

The convenience of the solution above is that you are now able to mock as many subsequent click prompts as you want, regardless of whether they are visible or hidden and in whichever order you arrange them, since mocking each type is handled non-monolithically.