如何为含input的异步函数编写unittest.mock测试？解决等待输入问题

问题描述 #

异步主函数start_script通过input获取用户输入分支执行逻辑，编写unittest测试时用unittest.mock.patch模拟输入，但运行测试时程序卡在控制台等待手动输入，模拟未生效。

原异步函数代码：

import asyncio

async def start_script():
    script = int(input("Введите:\n1 для запуска парсера видео\n2 для запуска загрузчика видео: \n"))
    if script == 1:
        number = int(input("Введите количество видео для парсинга: "))
        # await main.info_videos(number=number)
    elif script == 2:
        number_pc = int(input("Введите номер ПК: "))
        count_video = int(input("Количество повторных публикаций: "))
        print("Введите диапазон времени для публикации (в минутах)")
        number1 = int(input("От (в минутах): "))
        number2 = int(input("До (в минутах): "))
        await get_video_process(number1, number2, number_pc, count_video)
    else:
        print("Ошибка ввода!")
        await start_script()


asyncio.run(start_script())

原测试代码：

import unittest
from unittest.mock import patch

from web import start_script


class TestStartScript(unittest.TestCase):
    @patch("builtins.input", side_effect=["1", "5"])
    def test_start_script(self, mock_input):
        self.assertIsNone(start_script())


if __name__ == "__main__":
    unittest.main()

问题原因 #

1. 模块级代码自动执行：原异步函数所在模块的最后一行asyncio.run(start_script())未用if __name__ == "__main__"包裹，测试代码导入start_script时，这段代码会立即执行，此时测试的mock还未生效，导致程序卡在手动输入环节。
2. 异步函数未正确运行：测试中直接调用start_script()仅返回协程对象，未实际执行，后续即使进入测试逻辑，协程也不会运行。

解决方案 #

步骤1：修复模块级执行逻辑 #

修改原异步函数所在模块，将启动代码包裹在if __name__ == "__main__"中，避免导入时自动执行：

import asyncio

async def start_script():
    # 原函数逻辑不变
    script = int(input("Введите:\n1 для запуска парсера видео\n2 для запуска загрузчика видео: \n"))
    if script == 1:
        number = int(input("Введите количество видео для парсинга: "))
        # await main.info_videos(number=number)
    elif script == 2:
        number_pc = int(input("Введите номер ПК: "))
        count_video = int(input("Количество повторных публикаций: "))
        print("Введите диапазон времени для публикации (в минутах)")
        number1 = int(input("От (в минутах): "))
        number2 = int(input("До (в минутах): "))
        await get_video_process(number1, number2, number_pc, count_video)
    else:
        print("Ошибка ввода!")
        await start_script()

# 仅直接运行模块时才启动
if __name__ == "__main__":
    asyncio.run(start_script())

步骤2：修改测试代码，正确运行异步函数 #

使用asyncio.run()执行异步协程，同时mock依赖的异步函数避免实际执行：

import unittest
import asyncio
from unittest.mock import patch

from web import start_script


class TestStartScript(unittest.TestCase):
    @patch("builtins.input", side_effect=["1", "5"])
    def test_start_script_parser(self, mock_input):
        # 运行异步函数
        asyncio.run(start_script())
        # 验证输入调用符合预期
        mock_input.assert_any_call("Введите:\n1 для запуска парсера видео\n2 для запуска загрузчика видео: \n")
        mock_input.assert_any_call("Введите количество видео для парсинга: ")

    @patch("builtins.input", side_effect=["2", "1", "3", "10", "20"])
    @patch("web.get_video_process")  # 模拟异步函数，避免实际执行业务逻辑
    def test_start_script_downloader(self, mock_get_process, mock_input):
        asyncio.run(start_script())
        # 验证异步函数被正确调用
        mock_get_process.assert_awaited_once_with(10, 20, 1, 3)

    @patch("builtins.input", side_effect=["3", "1", "5"])
    @patch("builtins.print")
    def test_start_script_invalid_input(self, mock_print, mock_input):
        asyncio.run(start_script())
        # 验证错误提示被打印
        mock_print.assert_any_call("Ошибка ввода!")


if __name__ == "__main__":
    unittest.main()

关键要点 #

• 模块启动代码必须用if __name__ == "__main__"包裹，防止导入时自动执行
• 测试异步函数需通过asyncio.run()运行协程，否则协程不会实际执行
• 对于函数内部调用的其他异步方法，需mock隔离，避免测试依赖外部业务逻辑

内容的提问来源于stack exchange，提问作者Николай Сероух