How to test if name == "main" with passing command line arguments?

Question

Hi I want to test my executable module main.py. In this module there is function main() that takes two arguments:

# main.py

def main(population_size: int, number_of_iterations: int):
    ...

At the bottom of this module there is logic that takes command line arguments and executes main function:

# main.py

if __name__ == "__main__":
    # create parser and handle arguments
    PARSER = argparse.ArgumentParser()
    PARSER.add_argument("--populationSize",
                        type=int,
                        default=-1,
                        help="Number of individuals in one iteration")
    PARSER.add_argument("--numberOfIterations",
                        type=int,
                        default=-1,
                        help="Number of iterations in one run")
    # parse the arguments
    ARGS = PARSER.parse_args()

    main(ARGS.populationSize, ARGS.numberOfIterations)

I want to test passing command line arguments. My test method that doesn't work:

# test_main.py

@staticmethod
@mock.patch("argparse.ArgumentParser.parse_args")
@mock.patch("main.main")
def test_passing_arguments(mock_main, mock_argparse):
    """Test passing arguments."""
    mock_argparse.return_value = argparse.Namespace(
         populationSize=4, numberOfIterations=3)
    imp.load_source("__main__", "main.py")

    mock_main.assert_called_with(4, 3)

The error that I get is that mock_main is not called. I don't know why. To my understanding I mocked main function from main module. Mock of main function is neccessary becouse it's time consuming, and what I want to only test here is that parameters are passed correctly.

From this post I took way of mocking argparse module.

chepner · Answer 1 · 2020-02-13T14:14:46.957

Like all code you want to test, wrap it in a function.

def parse_my_args(argv=None):
    PARSER = argparse.ArgumentParser()
    PARSER.add_argument("--populationSize",
                    type=int,
                    default=-1,
                    help="Number of individuals in one iteration")
    PARSER.add_argument("--numberOfIterations",
                    type=int,
                    default=-1,
                    help="Number of iterations in one run")
    # parse the arguments
    return PARSER.parse_args(argv)


if __name__ == '__main__':
    args = parse_my_args()
    main(args.populationSize, args.numberOfIterations)

ArgumentParser.parse_args processes whatever list of strings you pass it. When you pass None, it uses sys.argv[1:] instead.

Now you can test parse_my_args simply by passing whatever list of arguments you want.

# test_main.py

@staticmethod
def test_passing_arguments():
    """Test passing arguments."""
    args = parse_my_args(["--populationSize", "4", "--numberOfIterations", "3"])
    assert args.populationSize == 4
    assert args.numberOfIterations == 3

If you further want to verify that the correct arguments are passed to main, wrap that in a function and use mock as you did above.

def entry_point(argv=None):
    args = parse_my_args(argv)
    main(args.populationSize, args.numberOfIterations)

if __name__ == '__main__':
    entry_point()

and

@staticmethod
@mock.patch("main.main")
def test_passing_arguments(mock_main):
    """Test passing arguments."""
    entry_point(["--populationSize", "4", "--numberOfIterations", "3"])
    mock_main.assert_called_with(4, 3)

score 3 · Answer 2 · answered Feb 13 '20 at 13:37

I usually write my command-line code like this. First rename your existing main function to something else, like run() (or whatever):

def run(population_size: int, number_of_iterations: int):
    ...

Then write a main() function which implements the command-line interface and argument parsing. Have it accept argv as an optional argument which is great for testing:

def main(argv=None):
    parser = argparse.ArgumentParser()
    ...
    args = parser.parse_args(argv)
    run(args.popuplation_size, args.number_of_iterations)

Then in the module body just put:

if __name__ == '__main__':
    sys.exit(main())

Now you have a proper main() function that you can easily test without fussing about the context in which it was called or doing any sort of weird monkeypatching, e.g. like:

main(['--populationSize', '4', '--numberOfIterations', '3'])

How to test if __name__ == "__main__" with passing command line arguments?

2 Answers2

How to test if name == "main" with passing command line arguments?