How to test or mock "if __name__ == '__main__'" contents

Say I have a module with the following:

def main():
    pass

if __name__ == "__main__":
    main()

I want to write a unit test for the bottom half (I'd like to achieve 100% coverage). I discovered the runpy builtin module that performs the import/__name__-setting mechanism, but I can't figure out how to mock or otherwise check that the main() function is called.

This is what I've tried so far:

import runpy
import mock

@mock.patch('foobar.main')
def test_main(self, main):
    runpy.run_module('foobar', run_name='__main__')
    main.assert_called_once_with()

Solution 1:

I will choose another alternative which is to exclude the if __name__ == '__main__' from the coverage report , of course you can do that only if you already have a test case for your main() function in your tests.

As for why I choose to exclude rather than writing a new test case for the whole script is because if as I stated you already have a test case for your main() function the fact that you add an other test case for the script (just for having a 100 % coverage) will be just a duplicated one.

For how to exclude the if __name__ == '__main__' you can write a coverage configuration file and add in the section report:

[report]

exclude_lines =
    if __name__ == .__main__.:

More info about the coverage configuration file can be found here.

Hope this can help.

Solution 2:

You can do this using the imp module rather than the import statement. The problem with the import statement is that the test for '__main__' runs as part of the import statement before you get a chance to assign to runpy.__name__.

For example, you could use imp.load_source() like so:

import imp
runpy = imp.load_source('__main__', '/path/to/runpy.py')

The first parameter is assigned to __name__ of the imported module.

Solution 3:

Whoa, I'm a little late to the party, but I recently ran into this issue and I think I came up with a better solution, so here it is...

I was working on a module that contained a dozen or so scripts all ending with this exact copypasta:

if __name__ == '__main__':
    if '--help' in sys.argv or '-h' in sys.argv:
        print(__doc__)
    else:
        sys.exit(main())

Not horrible, sure, but not testable either. My solution was to write a new function in one of my modules:

def run_script(name, doc, main):
    """Act like a script if we were invoked like a script."""
    if name == '__main__':
        if '--help' in sys.argv or '-h' in sys.argv:
            sys.stdout.write(doc)
        else:
            sys.exit(main())

and then place this gem at the end of each script file:

run_script(__name__, __doc__, main)

Technically, this function will be run unconditionally whether your script was imported as a module or ran as a script. This is ok however because the function doesn't actually do anything unless the script is being ran as a script. So code coverage sees the function runs and says "yes, 100% code coverage!" Meanwhile, I wrote three tests to cover the function itself:

@patch('mymodule.utils.sys')
def test_run_script_as_import(self, sysMock):
    """The run_script() func is a NOP when name != __main__."""
    mainMock = Mock()
    sysMock.argv = []
    run_script('some_module', 'docdocdoc', mainMock)
    self.assertEqual(mainMock.mock_calls, [])
    self.assertEqual(sysMock.exit.mock_calls, [])
    self.assertEqual(sysMock.stdout.write.mock_calls, [])

@patch('mymodule.utils.sys')
def test_run_script_as_script(self, sysMock):
    """Invoke main() when run as a script."""
    mainMock = Mock()
    sysMock.argv = []
    run_script('__main__', 'docdocdoc', mainMock)
    mainMock.assert_called_once_with()
    sysMock.exit.assert_called_once_with(mainMock())
    self.assertEqual(sysMock.stdout.write.mock_calls, [])

@patch('mymodule.utils.sys')
def test_run_script_with_help(self, sysMock):
    """Print help when the user asks for help."""
    mainMock = Mock()
    for h in ('-h', '--help'):
        sysMock.argv = [h]
        run_script('__main__', h*5, mainMock)
        self.assertEqual(mainMock.mock_calls, [])
        self.assertEqual(sysMock.exit.mock_calls, [])
        sysMock.stdout.write.assert_called_with(h*5)

Blam! Now you can write a testable main(), invoke it as a script, have 100% test coverage, and not need to ignore any code in your coverage report.

Solution 4:

Python 3 solution:

import os
from importlib.machinery import SourceFileLoader
from importlib.util import spec_from_loader, module_from_spec
from importlib import reload
from unittest import TestCase
from unittest.mock import MagicMock, patch
    

class TestIfNameEqMain(TestCase):
    def test_name_eq_main(self):
        loader = SourceFileLoader('__main__',
                                  os.path.join(os.path.dirname(os.path.dirname(__file__)),
                                               '__main__.py'))
        with self.assertRaises(SystemExit) as e:
            loader.exec_module(module_from_spec(spec_from_loader(loader.name, loader)))

Using the alternative solution of defining your own little function:

# module.py
def main():
    if __name__ == '__main__':
        return 'sweet'
    return 'child of mine'

You can test with:

# Override the `__name__` value in your module to '__main__'
with patch('module_name.__name__', '__main__'):
    import module_name
    self.assertEqual(module_name.main(), 'sweet')

with patch('module_name.__name__', 'anything else'):
    reload(module_name)
    del module_name
    import module_name
    self.assertEqual(module_name.main(), 'child of mine')