CODEWITHSUNDEEP
pythonpytestargparse
Dotan
Dotan

Reputation: 7622

Setting command line arguments for main function tests

I have a main() function in python that gets command line arguments. Is there a way for me to write pytest tests for this function and define the arguments in the code?

e.g.

def main():
    # argparse code
    args, other = arg_parser.parse_known_args()
    return args.first_arg


def test_main():
    res = main() # call with first_arg = "abc"
    assert(res == "abc")

Upvotes: 19

Views: 15116

Answers (4)

Little Bobby Tables
Little Bobby Tables

Reputation: 4744

For an entrypoint.py this might be particularly simple to add an extra argument. In this case --disable-warnings

import sys

import pytest


if __name__ == "__main__":
    exit_code = pytest.main(sys.argv[1:] + ["--disable-warnings"])

Upvotes: 0

leafmeal
leafmeal

Reputation: 2092

To add to the previous answers, instead of modifying sys.argv It is safer to use a context manager which can cover up and protect the underlying object. An example would be

with unittest.mock.patch('sys.argv', ['program_name', '--option1', 'inputFile']):
    main()

This works only with python3. For python2 the Mock library does the trick.

I found this solution on a different stackoverflow post here.

Upvotes: 19

hpaulj
hpaulj

Reputation: 231385

parse_args takes a argv parameter. The docs uses this repeatedly in it's examples

parser = argparse.ArgumentParser()
parser.add_argument('--foo', action='store_true')
parser.add_argument('bar')
parser.parse_known_args(['--foo', '--badger', 'BAR', 'spam'])

where the string list replicates sys.argv[1:] that it would get from the commandline. If the argument is None (or omitted) the parser uses sys.argv[1:].

So if 

def main(argv=None):
    # argparse code
    args, other = arg_parser.parse_known_args(argv)
    return args.first_arg

You could test with

main(['foo', '-f','v'])

The unittesting file for argparse.py uses both this approach, and your's of modifying sys.argv directly.

https://docs.python.org/3/library/argparse.html#beyond-sys-argv

https://docs.python.org/3/library/argparse.html#partial-parsing

Upvotes: 8

Dotan
Dotan

Reputation: 7622

The best solution I found so far is this

def test_main():
    sys.argv = ["some_name", "abc"]
    res = main()

and for flags:

sys.argv.append("-f")
sys.argv.append("v")

Upvotes: 5

Related Questions