python - 如何通过传递命令行参数测试__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函数的mock是必要的，因为它很耗时，我只想在这里测试的是参数是否正确传递。

来自 this post我采取了模拟 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，请将 that 包装在一个函数中并像上面那样使用 mock。

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)