how to test python file with pytest having if __name__ == '__main__': with arguments

Viewed 650

I want to test a python file with pytest which contains a if __name__ == '__main__': it also has arguments parsed in it. the code is something like this:

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Execute job.')

    parser.add_argument('--env', required=True, choices=['qa', 'staging', 'prod'])

    args = parser.parse_args()

    # some logic

The limitation here is I cannot add a main() method and wrap the logic in if __name__ == '__main__': inside it and run it from there! The code is a legacy code and cannot be changed!

I want to test this with pytest, I wonder how can I run this python file with some arguments inside my test?

2 Answers

Command line arguments are an input/output mechanism like any other. The key is to isolate it in a "boundary layer", and have your main program not depend on them directly.

In this case, rather than making your program access sys.argv directly (which is essentially a global variable), make it so your program is wrapped in an "entry point" function (e.g. "main") that takes args as an explicit parameter.

If you don't pass any arguments to ArgumentParser.parse_args, defaults to accessing sys.argv for you. Instead, just pass along your args param.

It might look something like:

def main(args):
    # some logic
    pass

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

Your unit tests can call this entry point function and pass in any args they want:

def test_main_does_foo_when_bar():
   result = main(["bar"])
   assert "foo" == result

Generally I make a main function which always takes arguments and call it with command line arguments when needed.

def main(args):
    # parser : argparse.ArgumentParser
    args = parser.parse_args(args)
    pass

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

When run from command line it will be called with command line arguments. It can be called at anytime by calling main directly.

Related