Pytest如何测试返回值依赖用户输入的Python函数

测试这类依赖用户终端输入的函数，不需要真的手动交互输入，用pytest内置的monkeypatch夹具替换掉Python内置的input方法，模拟输入值/抛出异常，就能覆盖所有逻辑分支。

合法输入返回场景 #

验证输入属于0-4的合法选项时，函数能正确返回对应整数：

def test_prompt_valid_choice(monkeypatch):
    # 单值测试
    monkeypatch.setattr("builtins.input", lambda _: "1")
    assert prompt() == 1

    # 全量合法值遍历测试
    for valid_opt in ["0", "1", "2", "3", "4"]:
        monkeypatch.setattr("builtins.input", lambda _: valid_opt)
        assert prompt() == int(valid_opt)

非法输入重试场景 #

验证输入非法值时函数会循环重问，直到收到合法输入才返回：

def test_prompt_retry_on_invalid(monkeypatch):
    # 按顺序返回模拟输入：前两次是非法值，第三次是合法值
    mock_input_sequence = iter(["wrong_input", "7", "4"])
    monkeypatch.setattr("builtins.input", lambda _: next(mock_input_sequence))
    assert prompt() == 4

异常终止场景 #

验证输入时触发EOFError、KeyboardInterrupt时，函数会按预期退出：

import pytest

def test_prompt_exit_on_eof(monkeypatch):
    def raise_eof(_):
        raise EOFError
    monkeypatch.setattr("builtins.input", raise_eof)
    with pytest.raises(SystemExit) as exit_info:
        prompt()
    assert "Program Terminated." in str(exit_info.value)

def test_prompt_exit_on_keyboard_interrupt(monkeypatch):
    def raise_keyboard_interrupt(_):
        raise KeyboardInterrupt
    monkeypatch.setattr("builtins.input", raise_keyboard_interrupt)
    with pytest.raises(SystemExit) as exit_info:
        prompt()
    assert "Program Terminated." in str(exit_info.value)

注意：自定义的mock input函数必须预留一个位置参数，因为原生input()被调用时会传入提示字符串作为入参，缺少参数会触发运行时报错。

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