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

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.

One approach is to run the modules as scripts (e.g. os.system(…)) and compare their stdout and stderr output to expected values.

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.

My solution is to use imp.load_source() and force an exception to be raised early in main() by not providing a required CLI argument, providing a malformed argument, setting paths in such a way that a required file is not found, etc.

import imp    
import os
import sys

def mainCond(testObj, srcFilePath, expectedExcType=SystemExit, cliArgsStr=''):
    sys.argv = [os.path.basename(srcFilePath)] + (
        [] if len(cliArgsStr) == 0 else cliArgsStr.split(' '))
    testObj.assertRaises(expectedExcType, imp.load_source, '__main__', srcFilePath)

Then in your test class you can use this function like this:

def testMain(self):
    mainCond(self, 'path/to/main.py', cliArgsStr='-d FailingArg')

I found this solution helpful. Works well if you use a function to keep all your script code.
The code will be handled as one code line. It doesn’t matter if the entire line was executed for coverage counter (though this is not what you would actually actually expect by 100% coverage)
The trick is also accepted pylint. 😉

if __name__ == '__main__': 
    main()

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')

I did not want to exclude the lines in question, so based on this explanation of a solution, I implemented a simplified version of the alternate answer given here…

1. I wrapped if __name__ == "__main__": in a function to make it easily testable, and then called that function to retain logic:

# myapp.module.py

def main():
    pass

def init():
    if __name__ == "__main__":
        main()

init()

2. I mocked the __name__ using unittest.mock to get at the lines in question:

from unittest.mock import patch, MagicMock
from myapp import module

def test_name_equals_main():
  # Arrange
  with patch.object(module, "main", MagicMock()) as mock_main:
    with patch.object(module, "__name__", "__main__"):
         # Act
         module.init()

  # Assert
  mock_main.assert_called_once()

If you are sending arguments into the mocked function, like so,

if __name__ == "__main__":
    main(main_args)

then you can use assert_called_once_with() for an even better test:

expected_args = ["expected_arg_1", "expected_arg_2"]
mock_main.assert_called_once_with(expected_args)

If desired, you can also add a return_value to the MagicMock() like so:

with patch.object(module, "main", MagicMock(return_value='foo')) as mock_main:

If it’s just to get the 100% and there is nothing "real" to test there, it is easier to ignore that line.

If you are using the regular coverage lib, you can just add a simple comment, and the line will be ignored in the coverage report.

if __name__ == '__main__':
    main()  # pragma: no cover

https://coverage.readthedocs.io/en/coverage-4.3.3/excluding.html

Another comment by @ Taylor Edmiston also mentions it

To import your "main" code in pytest in order to test it you can import main module like other functions thanks to native importlib package :

def test_main():
    import importlib
    loader = importlib.machinery.SourceFileLoader("__main__", "src/glue_jobs/move_data_with_resource_partitionning.py")
    runpy_main = loader.load_module()
    assert runpy_main()