如何测试或模拟"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__'覆盖率报告,当然,只有在测试中已经有main()函数的测试用例时才能这样做.

至于为什么我选择排除而不是为整个脚本编写一个新的测试用例是因为如果我说你已经为你的main()函数有一个测试用例,你为脚本添加了另一个测试用例(只是如果有100%的覆盖率,那将只是一个重复的.

有关如何排除,main()您可以编写coverage配置文件并添加到部分报告中:

[report]

exclude_lines =
    if __name__ == .__main__.:

有关coverage配置文件的更多信息,请访问此处.

希望这可以提供帮助.

Answer 2

您可以使用imp模块而不是import语句来执行此操作.该import语句的问题在于,'__main__'在您有机会分配之前,运行测试作为import语句的一部分runpy.__name__.

例如,您可以这样使用imp.load_source():

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

第一个参数分配给__name__导入的模块.

Answer 3

哇,我参加派对有点晚了,但我最近碰到了这个问题,我想我想出了一个更好的解决方案,所以这里......

我正在开发一个包含十几个脚本的模块,所有脚本都以这个精确的copypasta结尾:

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

然后将此gem放在每个脚本文件的末尾:

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)

布拉姆!现在,您可以编写一个testable main(),将其作为脚本调用,具有100%的测试覆盖率,而不需要忽略覆盖率报告中的任何代码.

Answer 4

Python 3 解决方案：

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

使用定义自己的小函数的替代解决方案：

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

您可以使用以下方法进行测试：

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