[Python 杂记] 使用unittest.mock模拟键盘输入

由Jeza Chen 发表于 August 2, 2022

在写Mumu-Notion的时候,由于测试代码有一些地方需要键盘输入(比如输入token等信息,以及输入回车继续下一个步骤),如果直接跑pytest的话会直接失败。

解决方法

可以使用unittest模块的unittest.mock.patch来为内置的input函数”打补丁“,以实现模拟键盘的输入。

mumu-notion的测试示例需要先输入Notion的Token, 然后输入一个页面的page id,最后疯狂按回车继续下一个步骤。因此,我们编写具体代码如下:

def example_code_input_generator(auth_token: str,
                                 base_page_id: str):
    """ A generator to mock input
    1. returns NOTION_AUTH_TOKEN
    2. returns NOTION_BASE_PAGE_ID
    3... always returns \n
    """
    yield auth_token
    yield base_page_id

    while True:
        yield '\n'

with unittest.mock.patch("builtins.input") as m:
    # mock the input
    m.side_effect = example_code_input_generator(auth_token, base_page_id)
    # ...下面为测试代码...

其中,我们在测试代码中为builtins.input打一个名为m的补丁,然后使用m.side_effect来模拟多次调用的返回数据(不能使用m.return_value,这个每次调用都只能返回同一个值)。example_code_input_generator生成器用于模拟具体的输入数据(可以使用list等数据结构,但由于我们的测试代码中按回车继续的次数不尽相同,因此使用生成器写一个while循环不断yield回车更好)。

这样,我们运行测试代码,就可以发现测试过程中自动实现键盘输入了!

参考资料

  1. Python的Mock文档

  2. Stack Overflow的问题