如何通过传递命令行参数测试__name__ ==“__main__”？

Posted 2023-03-25

【中文标题】如何通过传递命令行参数测试__name__ ==“__main__”？

【英文标题】：How to test if __name__ == "__main__" with passing command line arguments?

【发布时间】：2020-05-29 05:38:31

【问题描述】：

您好，我想测试我的可执行模块main.py。 在这个模块中有一个函数main()，它接受两个参数：

# main.py

def main(population_size: int, number_of_iterations: int):
    ...

在这个模块的底部有一个接受命令行参数并执行main函数的逻辑：

# main.py

if __name__ == "__main__":
    # create parser and handle arguments
    PARSER = argparse.ArgumentParser()
    PARSER.add_argument("--populationSize",
                        type=int,
                        default=-1,
                        help="Number of individuals in one iteration")
    PARSER.add_argument("--numberOfIterations",
                        type=int,
                        default=-1,
                        help="Number of iterations in one run")
    # parse the arguments
    ARGS = PARSER.parse_args()

    main(ARGS.populationSize, ARGS.numberOfIterations)

我想测试传递命令行参数。我的测试方法不起作用：

# test_main.py

@staticmethod
@mock.patch("argparse.ArgumentParser.parse_args")
@mock.patch("main.main")
def test_passing_arguments(mock_main, mock_argparse):
    """Test passing arguments."""
    mock_argparse.return_value = argparse.Namespace(
         populationSize=4, numberOfIterations=3)
    imp.load_source("__main__", "main.py")

    mock_main.assert_called_with(4, 3)

我得到的错误是 mock_main 没有被调用。我不知道为什么。据我了解，我从 main 模块中模拟了 main 函数。模拟main函数是必要的，因为它很耗时，我只想在这里测试的是参数是否正确传递。

来自this post 我采取了模拟 argparse 模块的方式。

【参考方案1】：

我通常这样编写命令行代码。首先将您现有的 main 函数重命名为其他名称，例如 run()（或其他）：

def run(population_size: int, number_of_iterations: int):
    ...

然后编写一个main() 函数来实现命令行接口和参数解析。让它接受argv 作为一个非常适合测试的可选参数：

def main(argv=None):
    parser = argparse.ArgumentParser()
    ...
    args = parser.parse_args(argv)
    run(args.popuplation_size, args.number_of_iterations)

然后在模块体中放：

if __name__ == '__main__':
    sys.exit(main())

现在你有了一个合适的 main() 函数，你可以轻松地测试它，而不用担心它被调用的上下文或做任何奇怪的猴子补丁，例如喜欢：

main(['--populationSize', '4', '--numberOfIterations', '3'])

【参考方案2】：

像所有要测试的代码一样，将其包装在一个函数中。

def parse_my_args(argv=None):
    PARSER = argparse.ArgumentParser()
    PARSER.add_argument("--populationSize",
                    type=int,
                    default=-1,
                    help="Number of individuals in one iteration")
    PARSER.add_argument("--numberOfIterations",
                    type=int,
                    default=-1,
                    help="Number of iterations in one run")
    # parse the arguments
    return PARSER.parse_args(argv)


if __name__ == '__main__':
    args = parse_my_args()
    main(args.populationSize, args.numberOfIterations)

ArgumentParser.parse_args 处理您传递给它的任何字符串列表。当你传递None 时，它会使用sys.argv[1:]。

现在您可以通过传递您想要的任何参数列表来测试parse_my_args。

# test_main.py

@staticmethod
def test_passing_arguments():
    """Test passing arguments."""
    args = parse_my_args(["--populationSize", "4", "--numberOfIterations", "3"])
    assert args.populationSize == 4
    assert args.numberOfIterations == 3

---

如果您想进一步验证是否将正确的参数传递给 main，请将 包装在一个函数中并像上面那样使用模拟。

def entry_point(argv=None):
    args = parse_my_args(argv)
    main(args.populationSize, args.numberOfIterations)

if __name__ == '__main__':
    entry_point()

和

@staticmethod
@mock.patch("main.main")
def test_passing_arguments(mock_main):
    """Test passing arguments."""
    entry_point(["--populationSize", "4", "--numberOfIterations", "3"])
    mock_main.assert_called_with(4, 3)