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

Question

假设我有一个具有以下内容的模块：

def main():
    pass

if __name__ == "__main__":
    main()

我想为下半部分编写一个单元测试（我想实现 100% 的覆盖率）。我发现了 runpy 内置模块执行 import/__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

我不想排除有问题的行，因此基于 这个解决方案的解释，我实现了此处给出的替代答案的简化版本...

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

# myapp.module.py

def main():
    pass

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

init()

2. 我使用 unittest 模拟了 __name__。模拟来获取有问题的行：

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_used_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 2

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

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

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

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

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

Answer 3

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

Answer 4

我发现这个解决方案很有帮助。如果您使用函数来保留所有脚本代码，则效果很好。
该代码将作为一行代码进行处理。整行是否针对覆盖率计数器执行并不重要（尽管这并不是您实际期望的 100% 覆盖率）
这个技巧也被 pylint 所接受。 ;-)

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

Answer 5

你的代码对我来说似乎大部分是正确的。由于我没有找到适合我的用例的答案，因此我在这里总结我的发现。我想要一些稍微不同的东西：使用 pytest 来测试也接受参数的脚本。
下面的代码应该适用于 python 3.1+，它使用 pytest 和 runpy，并使用 unittest 来传递参数。我没有使用 mock 因为我还不熟悉它。

测试脚本

输入脚本 totest.py

def main():
    print("Hello")

if __name__ == '__main__':
   main()

Pytest 脚本 test_totest.py

import runpy

def test_main():
    runpy.run_path('totest.py', run_name='__main__')

带参数的 Pytest 脚本

如果您使用 argparse 或 sys，这很有用脚本中的 .argv：

import unittest
import runpy

def test_main():
    with unittest.mock.patch('sys.argv', ["progname", "arg1", "arg2"]):
        runpy.run_path('totest.py', run_name='__main__')

测试模块

我也将其包含在此处以更直接地回答问题。

模块脚本 totest/__main__.py

def main():
    print("Hello")

if __name__ == '__main__':
   main()

Pytest 脚本 test_totest_module.py

这与之前的脚本几乎相同（run_module 而不是 run_path< /code>)：

import runpy

def test_main():
    runpy.run_module('totest.py', run_name='__main__')

检查 main() 函数是否被调用。

在上面的示例中，main() 函数始终被调用。如果您对此表示怀疑，请随意：

• 在 if 中调用 main() 后添加一个 assert False：您将看到“Hello “显示在 Pytest 的标准输出日志上（这对于我的示例来说有点特定）。
• 在 main 函数中添加一个 assert False：pytest 会让你的测试失败，证明 main 被执行了
• 如果你想以编程方式检查它，runpy返回包含模块全局变量的字典。因此，您可以在 main() 函数中更改其中一些模块全局变量，并稍后检查它们是否已更改。您还可以使用此机制在执行 main() 后检查模块或模块全局变量的状态，这对我来说比仅仅检查函数是否被执行更有意义。

Here's a full example of the latter approach

totest/__init__.py

class Testdata:  # This is defined here so that the class is present in the module globals
    main_was_executed = False

totest/__main__.py

from . import Testdata  # Import the previously-defined class
    
def main():
    print("Hello")
    Testdata.main_was_executed = True  # Change the state of the module global

if __name__ == '__main__':
    main()

Pytest脚本 test_totest_module.py

import runpy

def test_main():
    main = runpy.run_module('totest', run_name='__main__')  # Retrieve module globals
    assert main['Testdata'].main_was_executed  # Access the member of the module global

当然，你也可以定义自己的函数来断言这一点，将其导出到模块全局变量中，并从 pytest 脚本中调用它。

Answer 6

我的解决方案是使用 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 7

要在 pytest 中导入“ma​​in”代码以进行测试，您可以像其他函数一样导入主模块，这要归功于本机 importlib 包：

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

Answer 8

由于 imp 已被弃用，因此这里有一个 importlib 版本，以及命令行参数模拟。在exec_module之后，foobar将被完全初始化，并且foobar.main已经被执行。

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)

Answer 9

有人欣赏一点横向思维吗？

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

我希望这能引起笑声。

我个人不太关心覆盖工具。

但我实际上建议今后将其纳入我的项目中。令我烦恼的是，我的应用程序可能会在交付时缺少一对对其功能至关重要的行，但却通过了所有测试。

我欣然接受这并不完美（例如，可以在多行注释中找到匹配的字符串），但恕我直言，这总比没有好。

PS，您还可以免费获得 UTF-8 编码的基本奖励检查...

Answer 10

我将选择另一种选择，即从覆盖率报告中排除 if __name__ == '__main__' ，当然，只有当您已经有 main() 函数的测试用例时，您才可以这样做你的测试。

至于为什么我选择排除而不是为整个脚本编写一个新的测试用例，是因为如果正如我所说，您已经有一个针对您的 main() 函数的测试用例，那么您添加另一个测试用例脚本的测试用例（仅用于 100% 覆盖率）将只是一个重复的测试用例。

对于如何排除 if __name__ == '__main__' ，您可以编写一个覆盖率配置文件并在部分报告中添加：

[report]

exclude_lines =
    if __name__ == .__main__.:

有关覆盖率配置文件的更多信息可以找到此处。

希望这能有所帮助。

Answer 11

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

例如，您可以像这样使用 imp.load_source()：

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

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

Answer 12

哇，我参加聚会有点晚了，但我最近遇到了这个问题，我想我想出了一个更好的解决方案，所以这就是......

我正在开发一个包含十几个脚本的模块一切都以这个完全相同的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)

Blam！现在您可以编写一个可测试的 main()，将其作为脚本调用，拥有 100% 的测试覆盖率，并且不需要忽略覆盖率报告中的任何代码。

Answer 13

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