Unit testing __main__.py

Question

I have a Python package (Python 3.6, if it makes a difference) that I've designed to run as 'python -m package arguments' and I'd like to write unit tests for the __main__.py module. I specifically want to verify that it sets the exit code correctly. Is it possible to use runpy.run_module to execute my __main__.py and test the exit code? If so, how do I retrieve the exit code?

To be more clear, my __main__.py module is very simple. It just calls a function that has been extensively unit tested. But when I originally wrote __main__.py, I forgot to pass the result of that function to exit(), so I would like unit tests where the main function is mocked to make sure the exit code is set correctly. My unit test would look something like:

@patch('my_module.__main__.my_main', return_value=2)
def test_rc2(self, _):
    """Test that rc 2 is the exit code."""
    sys.argv = ['arg0', 'arg1', 'arg2', …]
    runpy.run_module('my_module')
    self.assertEqual(mod_rc, 2)

My question is, how would I get what I’ve written here as ‘mod_rc’?

Thanks.

Answer 1

Misko Hevery has said before (I believe it was in Clean Code Talks: Don't Look for Things but I may be wrong) that he doesn't know how to effectively unit test main methods, so his solution is to make them so simple that you can prove logically that they work if you assume the correctness of the (unit-tested) code that they call.

For example, if you have a discrete, tested unit for parsing command line arguments; a library that does the actual work; and a discrete, tested unit for rendering the completed work into output, then a main method that calls all three of those in sequence is assuredly going to work.

With that architecture, you can basically get by with just one big system test that is expected to produce something other than the "default" output and it'll either crash (because you wired it up improperly) or work (because it's wired up properly and all of the individual parts work).

---

At this point, I'm dropping all pretense of knowing what I'm talking about. There is almost assuredly a better way to do this, but frankly you could just write a shell script:

python -m package args
test $? -eq [expected exit code]

That will exit with error iff your program outputs incorrectly, which TravisCI or similar will regard as build failing.

Answer 2

__main__.py is still subject to normal __main__ global behavior — which is to say, you can implement your __main__.py like so

def main():
  # Your stuff

if __name__ == "__main__":
  main()

and then you can test your __main__ in whatever testing framework you like by using

from your_package.__main__ import main

As an aside, if you are using argparse, you will probably want:

def main(arg_strings=None):
  # …
  args = parser.parse_args(arg_strings)
  # …

if __name__ == "__main__":
  main()

and then you can override arg strings from a unit test simply with

from your_package.__main__ import main

def test_main():
  assert main(["x", "y", "z"]) == …

or similar idiom in you testing framework.