如何通过传递命令行参数测试 __name__ == "__main__" 是否？

Question

嗨，我想测试我的可执行模块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函数是必要的，因为它很耗时，我只想在这里测试参数是否正确传递。

从这篇文章中，我采用了模拟 argparse 模块的方法。

Answer 1

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

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)