如何测试接收用户输入并分支调用其他函数的Python函数？

我编写了如下display_menu函数，它接收用户输入，仅允许输入a或b，并根据输入分别调用new_user或existing_user函数。目前我已编写了部分unittest测试代码，但测试逻辑存在问题，请问如何完成该函数的完整测试？

函数代码 #

def display_menu() -> str:
    """
    Start menu that directs the user wether
    a new or existing option is chosen
    """
    clrscr()
    program_title()
    print("Hello and Welcome are you: \n")
    status = input("a)New User\nb)Exsisting user \n")
    while status not in ("a", "b"):
        print(Fore.RED + "Please choose between a or b")
        status = input("a)New User\nb)Existing user \n")
    if status == "a":
        new_user()
    elif status == "b":
        exsisting_user()

    return status

现有测试代码 #

import unittest
from unittest.mock import patch

from validate import display_menu,new_user

class TestValidate(unittest.TestCase):
    """
    testing  for user input   
    between two options offered
    """
    def test_a(self):
        input = __builtins__.input
         __builtins__.input = lambda _: 'a'
        self.assertTrue('please enter your username')

完整测试方案 #

你的现有测试存在几个问题：直接修改__builtins__.input不够规范，未验证函数核心逻辑（返回值、依赖函数调用），也没覆盖无效输入场景。以下是完整的测试实现：

核心思路 #

1. 使用unittest.mock.patch安全模拟用户输入、控制台输出，以及clrscr、program_title、new_user、existing_user等依赖函数，避免测试依赖外部环境。
2. 覆盖所有关键测试场景：有效输入a、有效输入b、先无效输入再输入有效值。

完整测试代码 #

import unittest
from unittest.mock import patch, MagicMock
from validate import display_menu

class TestValidate(unittest.TestCase):
    def setUp(self):
        # 统一模拟无需验证的依赖函数
        self.patcher_clrscr = patch('validate.clrscr').start()
        self.patcher_title = patch('validate.program_title').start()
        self.patcher_fore = patch('validate.Fore').start()
        # 模拟业务函数，用于验证调用情况
        self.mock_new_user = patch('validate.new_user').start()
        self.mock_existing_user = patch('validate.exsisting_user').start()

    def tearDown(self):
        # 清理所有mock资源
        patch.stopall()

    def test_input_a_returns_a_and_calls_new_user(self):
        # 模拟用户输入'a'
        with patch('builtins.input', return_value='a'):
            result = display_menu()
            # 验证返回值正确
            self.assertEqual(result, 'a')
            # 验证new_user被调用一次
            self.mock_new_user.assert_called_once()
            # 验证existing_user未被调用
            self.mock_existing_user.assert_not_called()

    def test_input_b_returns_b_and_calls_existing_user(self):
        # 模拟用户输入'b'
        with patch('builtins.input', return_value='b'):
            result = display_menu()
            self.assertEqual(result, 'b')
            self.mock_existing_user.assert_called_once()
            self.mock_new_user.assert_not_called()

    def test_invalid_input_then_valid_a(self):
        # 模拟先输入无效值'c'，再输入有效值'a'
        with patch('builtins.input', side_effect=['c', 'a']):
            # 捕获控制台输出，验证错误提示
            with patch('builtins.print') as mock_print:
                result = display_menu()
                self.assertEqual(result, 'a')
                # 验证错误提示被打印
                error_msg = self.patcher_fore.RED + "Please choose between a or b"
                mock_print.assert_any_call(error_msg)
                self.mock_new_user.assert_called_once()

代码说明 #

• setUp和tearDown：统一管理mock资源，保证每个测试用例独立且无重复代码。
• test_input_a_returns_a_and_calls_new_user：验证输入a时，函数返回a且正确调用new_user。
• test_input_b_returns_b_and_calls_existing_user：验证输入b时，函数返回b且正确调用exsisting_user（注意原函数中拼写为exsisting_user，测试时需对应）。
• test_invalid_input_then_valid_a：验证无效输入触发循环提示，后续输入有效值后流程正常，同时验证错误提示被输出。

内容的提问来源于stack exchange，提问作者user18816965