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 b

Answer 1

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.