How to write unit test for code with inputs (Python)

Question

I'm new to unit testing and I have to add a unit test class to a solution to a problem. The problem is the one found herehttps://practice.geeksforgeeks.org/problems/penalty-shooters/0# and my code is as follows. I'm having trouble figuring out how to unit test with user input.

T = int(input())

for x in range(T):
    # Allows us to input and apply method to all iterables in a | no. of tests \n x_energy y_energy z_energy | format. T = number of tests, x is Lohia energy, y is Gosu energy and z is Goalkeeper energy
    x, y, z = list(map(int, input().split(" ")))
    Goal_Lohia, Goal_Gosu = 0, 0

    # Session ends when goalkeepers energy decrements to 1
    while z > 1:
        # Goalkeepers energy in this case is a factor of both strikers' energy, so both score.
        if x % z == 0 and y%z == 0:
            x = x - 1
            Goal_Lohia = Goal_Lohia + 1
            y = y - 1
            Goal_Gosu = Goal_Gosu + 1
        # In this case, the goalkeepers energy is only a factor of Lohia's. Note Lohia goes first as the problem states.
        elif x % z == 0:
            x = x - 1
            Goal_Lohia = Goal_Lohia + 1
        # And in this case, it is only a factor of Gosu's
        elif y % z == 0:
            y = y - 1
            Goal_Gosu = Goal_Gosu + 1
        # If goalkeepers energy is a factor of neither Gosu nor Lohia, the goalkeeper will save (no goal scored) and will lose energy.
        else:
            z = z - 1
    print(Goal_Lohia, Goal_Gosu)

Many thanks!

Answer 1

You want to refactor your code so that you can call the code both by taking input from the user and by providing a hard-coded value. Then you can call one version when the program runs normally, and the other when running tests.

You can similarly separate out what is done with the output, so that you can print the output in the normal case, but check it for correctness in the testing case.

Here's what this could look like for your code:

def function_to_test(T):
    for x in range(T):
        # Allows us to input and apply method to all iterables in a | no. of tests \n x_energy y_energy z_energy | format. T = number of tests, x is Lohia energy, y is Gosu energy and z is Goalkeeper energy
        x, y, z = list(map(int, input().split(" ")))
        Goal_Lohia, Goal_Gosu = 0, 0
        ...
        return Goal_Lohia, Goal_Gosu

def function_that_takes_user__input():
    T = int(input())
    Goal_Lohia, Goal_Gosu = function_to_test(T)
    print(Goal_Lohia, Goal_Gosu)

def test_the_function():
    Goal_Lohia, Goal_Gosu = function_to_test(<test input 1>)
    # add code to test result here

    Goal_Lohia, Goal_Gosu = function_to_test(<test input 2>)
    # add code to test result here

Answer 2

The best way is to extract this logic as a pure function and pass the parsed data to it:

def logic(records):
    for x, y, z in records:
        ...

    return goal_lohia, goal_gosu
    

def main():
    t = int(input())
    records = []
    for _ in range(t):
        records.append(list(int(c) for c in input().split()))

    print(logic(records))

Now you can test logic() in isolation from the input/output concerns.