如何测试或模拟“if __name__ == '__main__'”中的内容

Question

假设我有一个模块，内容如下：

def main():
    pass

if __name__ == "__main__":
    main()

我想为下面的部分写一个单元测试（我希望能覆盖到100%）。我发现了一个叫做runpy的内置模块，它可以处理导入和__name__的设置机制，但我不知道怎么去模拟或者检查main()函数是否被调用。

这是我到目前为止尝试过的：

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

Answer 1

哇，我来得有点晚，不过我最近遇到了这个问题，觉得我想出了一个更好的解决办法，所以就分享给大家...

我在做一个模块，里面有十几个脚本，最后都以一段完全一样的代码结束：

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

这段代码虽然不算糟糕，但也没法测试。我的解决办法是在我的一个模块里写了一个新函数：

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

然后把这个小宝贝放在每个脚本文件的最后：

run_script(__name__, __doc__, main)

从技术上讲，这个函数无论你的脚本是作为模块导入还是直接运行，都会被执行。这没关系，因为这个函数实际上只有在脚本被运行时才会做事情。所以代码覆盖率工具会看到这个函数被运行，就会说“是的，100%代码覆盖率！”与此同时，我还写了三个测试来覆盖这个函数本身：

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

太棒了！现在你可以写一个可测试的 main() 函数，作为脚本运行，拥有100%的测试覆盖率，而且在你的覆盖率报告中不需要忽略任何代码。

Answer 2

你可以使用 imp 模块来实现这个，而不是用 import 语句。使用 import 语句的问题在于，检查 '__main__' 的那部分代码会在你有机会给 runpy.__name__ 赋值之前就执行了。

比如，你可以这样使用 imp.load_source()：

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

第一个参数会被赋值给被导入模块的 __name__。

Answer 3

我会选择另一种方法，就是把if __name__ == '__main__'从覆盖率报告中排除。当然，你只有在测试中已经有了对main()函数的测试用例时，才能这样做。

我选择排除而不是为整个脚本写一个新的测试用例，是因为如果你已经有了对main()函数的测试用例，那么为了达到100%的覆盖率而再添加一个测试用例其实就是重复的。

要排除if __name__ == '__main__'，你可以写一个覆盖率配置文件，并在报告部分添加：

[report]

exclude_lines =
    if __name__ == .__main__.:

关于覆盖率配置文件的更多信息可以在这里找到。

希望这能帮到你。