如何测试或模拟“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％覆盖率）将只是一个重复的。

有关如何排除if __name__ == '__main__'，您可以编写覆盖配置文件并添加部分报告：

[report]

exclude_lines =
    if __name__ == .__main__.:

可以找到有关coverage配置文件的更多信息here。

希望这可以提供帮助。

Answer 2

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

例如，您可以使用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)

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

Answer 4

一种方法是将模块作为脚本运行（例如os.system（...）），并将它们的stdout和stderr输出与预期值进行比较。

Answer 5

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

Answer 6

我不想排除有问题的行，因此基于 this explanation of a solution，我实现了简化版本的 the alternate answer given here...

1. 我将 if __name__ == "__main__": 包装在一个函数中以使其易于测试，然后调用该函数以保留逻辑：

# myapp.module.py

def main():
    pass

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

init()

2. 我使用 __name__ 嘲笑 unittest.mock 以获取有问题的行：

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 __name__ == "__main__":
    main(main_args)

然后您可以使用 assert_called_once_with() 进行更好的测试：

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

如果需要，您还可以像这样将 return_value 添加到 MagicMock()：

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

Answer 7

我的解决方案是使用imp.load_source()并通过不提供必需的CLI参数强制在main()中提出异常，提供格式错误的参数，以所需文件的方式设置路径找不到等等。

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)

然后在你的测试类中你可以像这样使用这个函数：

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

Answer 8

如果只是为了获得 100% 而没有任何“真实”可以测试，那么忽略该行会更容易。

如果您使用的是常规覆盖库，您只需添加一个简单的注释，该行将在覆盖报告中被忽略。

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

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

@Taylor Edmiston 的另一条评论也提到了它

Answer 9

我发现此解决方案很有帮助。如果使用函数保留所有脚本代码，效果很好。 该代码将作为一条代码行处理。覆盖范围计数器是否执行了整行都无关紧要（尽管这并不是100％覆盖率实际期望的值） 绝招也被接受。 ;-）

if __name__ == '__main__': \
    main()