如何测试或模拟“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()

- Nikolaj

Answer 1

我将选择另一种方法，即从覆盖报告中排除if __name__ == '__main__'，当然，您只能在测试用例中已经存在main()函数的情况下这样做。至于为什么我选择排除而不是为整个脚本编写新的测试用例，原因是如果如我所述，您已经有了一个main()函数的测试用例，那么添加一个脚本的其他测试用例（仅为了实现100％覆盖）就会成为重复的测试用例。

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

[report]

exclude_lines =
    if __name__ == .__main__.:

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

希望这可以帮助您。

- mouad

---

嗨，我添加了一个新答案，它提供了100％的测试覆盖率（带有测试！）并且不需要忽略任何内容。让我知道你的想法：https://dev59.com/pW025IYBdhLWcg3wsIMU#27084447 谢谢。 - robru

对于那些想知道的人：nose-cov 在底层使用了 coverage.py，因此具有上述内容的 .coveragerc 文件将完美地工作。 - Joscha

在我看来，即使我认为这篇回答很有趣且有用，它实际上并没有回答 OP 的问题。他想测试 main 函数是否被调用，而不是跳过这个检查。否则，在启动时，脚本实际上可以执行除了实际期望的内容之外的所有操作，而测试会显示“没问题，一切正常！”。即使实际上从未被调用，主函数也可以进行完全的单元测试。 - iacopo

虽然它可能不能回答OP的问题，但从实际角度来看，这是一个很好的答案，这也是我找到这个问题的原因。另一个类似的解决方案是使用 # pragma: no cover，就像这样 if __name__ == '__main__': # pragma: no cover。个人认为这种方法会让代码变得混乱且难看，所以我认为mouad的答案是最好的解决方案，但其他人可能会发现它有用。 - Taylor D. Edmiston

如果我们非常具体的话，我认为正则表达式应该使用['"]而不是.，像这样：__name__ == ['"]__main__['"]:。 - Taylor D. Edmiston

Answer 2

你可以使用imp模块而不是import语句来实现这一点。 import语句的问题在于，在你有机会为runpy.__name__赋值之前，'__main__'的测试将作为导入语句的一部分运行。

例如，你可以像这样使用imp.load_source():

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

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

- David Heffernan

---

imp 模块似乎与我在问题中使用的 runpy 模块非常相似。问题在于，模拟对象似乎无法在模块加载后和代码运行之前插入。你对此有什么建议吗？ - Nikolaj

imp模块现在已经被弃用，我尝试使用最新的importlib方法编辑您的答案，但是Stack Overflow说您的答案已经达到了编辑队列限制。 - iD_Sgh

Answer 3

哇，我来晚了，但我最近遇到了这个问题，并且我认为我想出了一个更好的解决方案，所以在这里分享一下...

我正在开发一个包含十几个脚本的模块，所有脚本都以相同的复制粘贴方式结尾：

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%测试覆盖率，并且不需要忽略你覆盖报告中的任何代码。

- robru

---

我欣赏你在寻找解决方案时的创造力和毅力，但如果你在我的团队里，我会否决这种编码方式。Python 的优势之一是它高度惯用。if __name__ == ... 是让模块脚本运行的唯一方法。任何 Python 程序员都会认识到这行代码并理解它的作用。你的解决方案只是为了满足一时的兴趣而使显而易见的事情变得混乱不清。就像我所说的：这是一个聪明的解决方案，但 聪明 并不总是等于 _正确_。 - mac

如果在以 if __name__ ... 缩进的块中有任何逻辑，那么你做错了，应该进行重构。在 if __name__... 下的唯一代码行应该是：main()。 - mac

@mac 我不确定我同意这个观点。是的，如果你有逻辑，你应该进行重构。但这并不意味着在 if __name__ ... 下面唯一可以有的东西就是 main()。例如，我喜欢使用argeparse并在 if __name__ ... 部分构建我的解析器。然后将我的主函数抽象出来，使用显式参数而不是像 main(parser.parse_args()) 这样的东西。这使得在需要时更容易从另一个模块调用 main()。否则，您必须构造一个 argeparse.Namespace() 对象并正确获取所有默认参数。或者有更惯用的方法吗？ - Michael Leonard

@MichaelLeonard - 我不确定我是否正确理解了你的问题。main是按照惯例应该在将模块作为脚本调用时运行的函数，因此它是解析代码的常规位置。如果您有一个要从模块内部公开的单个函数，则不应将其称为main，而应该命名为其他名称，并且main函数应调用它并传递解析的参数。或者我完全误解了你的问题？ - mac

在所有其他回复中，这个是唯一一个可以进行代码覆盖率而不会干扰Python内部和/或混淆导入路径的方法（imp模块需要绝对路径，不方便在tests/文件夹内使用）。这是我的版本（带有退出代码断言：https://gist.github.com/amarao/5d0d4346f9ce6183d05bdae197ba1fc4） - George Shuklin

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

- Samuel Marks

---

我有点困惑。你在这里提出了多个解决方案吗？我尝试了第一段代码，得到了“失败：未引发<class 'SystemExit'>”的结果。关于下一段代码：通常认为不要随意修改应用程序代码中的标准两行代码“if name == 'main': main()”是至关重要的。我不确定你是否在你的应用程序代码中改变了这个。 - mike rodent

Answer 5

我不想排除这些代码行，所以根据此解释提供的一种解决方案，我实现了这个备选答案的简化版本...

1. 我将 if __name__ == "__main__": 包装在一个函数中，以便轻松进行测试，然后调用该函数以保留逻辑：

# myapp.module.py

def main():
    pass

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

init()

我使用unittest.mock模块来模拟__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_called_once_with()进行更好的测试：

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

如果需要，你也可以像下面这样在 MagicMock() 中添加一个return_value:

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

- LHM

Answer 6

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

- Mr Fooz

---

在子进程中运行脚本并期望coverage.py跟踪执行的代码行并不像听起来那么容易，有关使此解决方案正常工作的更多信息可以在此处找到：http://nedbatchelder.com/code/coverage/subprocess.html - mouad

Answer 7

如果只是为了获得100％的覆盖率，而没有任何“真实”的测试需要进行，那么忽略该行会更容易。

如果您正在使用常规的覆盖库，您可以添加一个简单的注释，这样该行就会在覆盖报告中被忽略。

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

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

另外一条由@ Taylor Edmiston 发表的评论也提到了它

- mgershen

Answer 8

我发现这个解决方案很有帮助。如果你使用一个函数来保留所有的脚本代码，它会很好地运行。 该代码将被处理为一行代码。不用担心整行是否被执行以进行覆盖计数（虽然实际上这并不是你所期望的100%覆盖率）。 这个技巧也被pylint接受。 ;-)

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

- Stephan

Answer 9

我的解决方案是使用imp.load_source()，并在main()中提前引发异常，方法是不提供必需的CLI参数，提供格式不正确的参数，以这种方式设置路径，使得所需的文件无法找到等。

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

- polsar

Answer 10

为了测试你的“main”代码，可以使用本地的importlib包像导入其他函数一样导入main模块：

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

- iD_Sgh

---

这看起来非常聪明（我只对importlib的基本知识有所了解）。但是我仍然得到了"TypeError: 'module' object is not callable"的错误，尽管在SourceFileLoader构造函数中似乎已经正确替换了参数... - mike rodent