Testing User Input - Python

Question

I am having troubles with testing my input for my code in Python. I tried a couple of solutions, but there is something that I am missing, so I would appreciate it if you could give me some tips.

First here is a snippet from my main file of code that I want to test:

if __name__ == '__main__':

    n = int(input())
    m = int(input())

    grid = []

    for _ in range(n):
        grid.append(list(map(str, input().rstrip().split())))

    calculate(grid)

When I run my code, I input "n", then "m", then a grid is created according to the user input (each row on a new row..), and a function is executed that calculates something on the grid and the function returns the result. It all works great, but now I need to create a couple of test cases for it (that test different inputs against expected outputs).

First, I tried this: (on a separate .py file)

from unittest import mock
from unittest import TestCase
import main_file

class DictCreateTests(TestCase):
    @mock.patch('main_file.input', create=True)
    def testdictCreateSimple(self, mocked_input):
        mocked_input.side_effect = ['2', '2', 'R G B\nR G B'] #this is the input I need for my color grid
        self.assertEqual(calculate(grid), 2)

if __name__ == '__main__':
    unittest.main()

I then researched for some more options and I tried this option, which got me the closest:

import unittest
import os

class Test1(unittest.TestCase):

    def test_case1(self):
        input = "2\n2\nR G B\nR G B"
        expected_output = '2'
        with os.popen("echo " + input + "' | python main_file.py") as o:
            output = o.read()
        output = output.strip() # Remove leading spaces and LFs
        self.assertEqual(output, expected_output)

if __name__ == '__main__':
    unittest.main()

Unfortunately, even though it passed the test, I discovered that it always accepts the first letter/number of input as a result, when it compares it out to the expected output. So, I am thinking it has something to do with the multiple values I need to input. I tried separating them on different inputs (input1 + input2+ input3), but it still didn't work.

I would very much appreciate it if anyone could give me some tips on how to do it! Thank you in advance!

Answer 1

I suggest to refactor the code, so that you can test a function:

def create_grid_and_calculate(n, m):
    grid = []

    for _ in range(n):
        grid.append(list(map(str, input().rstrip().split())))

    return calculate(grid)


if __name__ == '__main__':

    n = int(input())
    m = int(input())

    create_grid_and_calculate(n, m)

Then

import unittest
import os
from main_file import create_grid_and_calculate

class Test1(unittest.TestCase):

    def test_case1(self):
        expected_output = '2'
        self.assertEqual(create_grid_and_calculate(2, 2), expected_output)
        self.assertEqual(create_grid_and_calculate(int("R G B"), int("R G B")), expected_output)

if __name__ == '__main__':
    unittest.main()

You could also replace your inputs with arguments passed on the command line, parsed with a dedicated module (argparse is standard for instance), to have a better control on your inputs.

import argparse

def create_grid_and_calculate(n, m):
    ...


def main(argv: list = None):
    parser = argparse.ArgumentParser(description="My script...")
    parser.add_argument(
        "-m",
        dest="m",
        action="store",
        type=int,
        help="parameter m",
    )
    parser.add_argument(
        "-n",
        dest="n",
        action="store",
        type=int,
        help="parameter n",
    )
    args = parser.parse_args(argv or [])

    create_grid_and_calculate(args.n, args.m)

if __name__ == '__main__':
    import sys
    sys.exit(main(sys.argv[1:]))

So you can also test the main function with different inputs (int, strings...).

Finally, pytest is a great unitary test framework build on top of unittest, maybe you can have a look.

---

EDIT: To define your grid, you do not need to input the sizes (n and m), and you do not need to input each row separately. Choose 1 separator for line (here the comma), another for columns (here the space), and you get:

import argparse

def main(argv: list = None):
    parser = argparse.ArgumentParser(description="My script...")
    parser.add_argument(
        "-g",
        "--grid",
        dest="grid",
        action="store",
        type=str,
        help="The grid, defined as 'x11 x12 ... x1n, x21 x22 ... x2n, ...'",
    )
    args = parser.parse_args(argv or [])

    # first we split on comma, then on space
    grid = [x.split() for x in args.grid.split(',')]
    print(grid)
    calculate(grid)

if __name__ == '__main__':
    import sys
    sys.exit(main(sys.argv[1:]))

and you run it like that: main_file.py -g '1 2 3, 4 5 6, 7 8 9'

[['1', '2', '3'], ['4', '5', '6'], ['7', '8', '9']]

To get a list of list of int, use:

grid = [[int(val) for val in row.split()] for row in args.grid.split(',')]

or perhaps more clearly:

grid = []
for row in args.grid.split(','):
    grid.append([])
    for val in row.split():
        grid[-1].append(int(val))