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

Question

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

Guenther

Answer 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.

Heeled

Heya, I've added a new answer that gives 100% test coverage (with tests !) and doesn't require ignoring anything. Let me know what you think: https://mcmap.net/q/216356/-how-to-test-or-mock-quot-if-__name__-39-__main__-39-quot-contents Thanks. – Maritamaritain

For those wondering: nose-cov uses coverage.py underneath, so a .coveragerc file with the above content will work just fine. – Banna

IMHO, even if I found it interesting and useful, this answer does not actually give a response to the OP. He want to test that main is called, not to skip this check. Otherwise, the script could actually do everything except what actually expected, when launched, with tests saying "OK, everything works!". And the main function could be fully unit-tested, even if being never called actually. – Chap

It might not give a response to OP, but it is a good answer for practical purposes which is how I found this question at least. A similar solution is using # pragma: no cover like so if __name__ == '__main__': # pragma: no cover. Personally I'm not willing to do this because it clutters the code and is pretty ugly, so I think mouad's answer is the best solution, but others may find it useful. – Approximation

@Heeled If we're being very specific, I think technically the regex line should use ['"] instead of . like: __name__ == ['"]__main__['"]:. – Approximation

Answer 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.

Sidras

The imp module seems to work much like the runpy module I used in the question. The problem is that the mock cannot (apparently) be inserted after the module was loaded and before the code was run. Do you have any suggestions for this? – Guenther

imp module is deprecated now, I tried to edit your answer to the newest way replacing it with importlib but stack overflow says there is an edit queue limit on your answer. – Tacy

Answer 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.

Maritamaritain

I appreciate the creativity and perseverance in finding a solution, but if you were in my team, I would veto this way of coding. One of the strenghts of Python is its being highly idiomatic. if __name__ == ... is the way to let a module script. Any pythonista will recognise that line and understand what it does. Your solution it's just obfuscating the obvious for no good reason other than scratching an intellectual itch. As I said: a clever solution, but clever does not always equate to correct. – Shalom

That's fine if you just have one module, or if each module does somthing different when called as a script, but as I said I had a dozen files with completely identical if __name__ == ... blocks at the end, which is a huge violation of Don't Repeat Yourself and also makes it difficult to fix bugs when you need to fix it identically in so many places. Unifying the logic like this increases testability and reduces the potential for bugs. If you're concerned about people not understanding it, name the function if_name_equals_main() and people will figure it out. – Maritamaritain

If you have any logic in the block indented under if __name__ ... then you are doing it wrong and should refactor. The only line of code under if __name__... should read: main(). – Shalom

@Shalom I don't know that I agree with that. Yes, if you have logic you should refactor. But that doesn't mean that the only thing you can have under if __name__ ... is main(). For example, I like to use argeparse and construct my parser in the if __name__ ... portion. Then abstract my main to use explicit args rather than something like: main(parser.parse_args()). This makes it easier to call main() from another module if needed. Otherwise you have to construct an argeparse.Namespace() object and get all the default args correct. Or is there a more idiomatic way to do it? – Vingtetun

@MichaelLeonard - I am not sure I understand your question correctly. main is - by convention - the function that should run when invoking the module as a script, so it is the conventional place where parsing code should go. If you have a single function you want to expose from within the module, that should not be called main but something else, and the main function should in turn invoke it passing the parsed arguments. Or am I misunderstanding your question entirely? – Shalom

I agree with mac on this one, main() should do your arg parsing and whatever function you're currently calling main should be named something else. – Maritamaritain

Out of all other replies this one is the single to do code coverage without messing with python internals and/or mangling import paths (imp module requires absolute path and is not convenient for use inside tests/ folder). Here my version (with exit code assertion: gist.github.com/amarao/5d0d4346f9ce6183d05bdae197ba1fc4) – Barbie

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

Conveyancer

I'm a bit confused. Are you proposing more than one solution here? I tried the first bit of code and got " Failed: DID NOT RAISE <class 'SystemExit'>". Re the next bit here: it is generally felt to be essential not to mess around with the standard pair of lines in the app code "if __name__ == '__main__': main()". I'm not sure whether or not you are changing this in your app code... – Ketchup

Answer 5

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:

Ferriage

Answer 6

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

Toxicant

Answer 7

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

Thunderstruck

Running the script in a sub process and expecting coverage.py to track the line executed is not at easy as it sound , more info to make this solution work can be found here: nedbatchelder.com/code/coverage/subprocess.html – Heeled

Answer 8

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

Marlenamarlene

Answer 9

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

Lylelyles

Answer 10

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

Tacy

This looks very clever (I only know the basics about importlib). But I just get "TypeError: 'module' object is not callable", despite seemingly substituting in the right parameters in the SourceFileLoader constructor... – Ketchup

Answer 11

Anyone appreciate a bit of lateral thinking?

def test_if_file_name_main_main_function_is_called():
    main_file_path = pathlib.Path.cwd().joinpath('__main__.py')
    with main_file_path.open(mode='r', encoding='utf-8') as f:
        content = f.read()
    pattern = '.*\nif\s+__name__\s*==\s*[\'"]__main__[\'"]\s*:\s*\n\s+main\(\)\s*\n.*'
    assert re.fullmatch(pattern, content, flags=re.DOTALL) != None

I hope that raises a laugh.

I'm personally not bothered much about coverage tools.

But I'm actually proposing to include this henceforth in my projects. It bugs me that my app could potentially be shipped with a pair of lines missing which are essential to its functioning, but pass all tests.

I readily accept that this ain't perfect (for example, the matching string could be found inside a multi-line comment), but it's better than nothing, IMHO.

PS you also get a basic bonus check on UTF-8 encoding thrown in for free...

Ketchup

Answer 12

As imp is deprecated, here is a importlib version, along with commandline arguments mock. After the exec_module, foobar would have been fully initiallized, and foobar.main has already been exectued.

import sys
class MyTest(unittest.TestCase):
    def test_main(self):
        import importlib.util
        from unittest.mock import patch

        testargs = ['foobar.py']
        with patch.object(sys, 'argv', testargs):
            spec = importlib.util.spec_from_file_location('__main__', 'foobar.py')
            runpy = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(runpy)
            # check result of foobar script in unittest from here
            self.assertIsNotNone(runpy)

Guy