如何在pytest中用monkeypatch实现mock input()的装饰器？

需求背景 #

需要在pytest中快速模拟内置input()函数，用可变字符串列表的迭代器替代它。现有实现存在两个痛点：

• 必须固定使用input_string_list作为参数名，灵活性不足
• 将逻辑拆分到其他文件时，需要同时导入多个对象，使用繁琐

现有可运行但有局限的实现 #

from typing import Callable
import pytest

def _create_patched_input(str_list: list[str]) -> Callable:
    str_iter = iter(str_list.copy())

    def patched_input(prompt: str) -> str:  # 和input函数签名一致
        val = next(str_iter)
        print(prompt + val, end="\n"),
        return val

    return patched_input

@pytest.fixture
def _mock_input(monkeypatch, input_string_list: list[str]):
    patched_input = _create_patched_input(input_string_list)
    monkeypatch.setattr("builtins.input", patched_input)

def mock_input(f):
    return pytest.mark.usefixtures("_mock_input")(f)

# 测试示例
def get_name(prompt: str) -> str:
    return input(prompt)

@mock_input
@pytest.mark.parametrize(
    "input_string_list",
    (["Alice", "Bob", "Carol"], ["Dale", "Evie", "Frank", "George"]),
)
def test_get_name(input_string_list):
    for name in input_string_list:
        assert get_name("What is your name?") == name

最优实现方案 #

核心思路 #

实现一个装饰器工厂，同时支持两种使用场景：

1. 直接传入固定字符串列表，快速生成mock
2. 配合@pytest.mark.parametrize使用自定义参数名，摆脱固定参数名限制

同时只需导入单个对象，解决多对象导入的繁琐问题。

完整实现代码 #

from typing import Callable, List, Union
import pytest

def _create_patched_input(str_list: List[str]) -> Callable:
    str_iter = iter(str_list.copy())

    def patched_input(prompt: str) -> str:
        val = next(str_iter)
        print(prompt + val, end="\n")
        return val

    return patched_input

def mock_input(input_source: Union[str, List[str]]):
    """
    装饰器工厂：用于模拟内置input()函数
    支持两种用法：
    1. 直接传入字符串列表：@mock_input(["Alice", "Bob"])
    2. 配合parametrize传入参数名：@mock_input("自定义参数名")
    """
    def decorator(test_func: Callable):
        if isinstance(input_source, str):
            # 动态创建fixture，通过参数名获取输入列表
            @pytest.fixture(autouse=True)
            def _dynamic_mock_input(monkeypatch, request):
                input_list = request.getfixturevalue(input_source)
                patched_input = _create_patched_input(input_list)
                monkeypatch.setattr("builtins.input", patched_input)
            
            return pytest.mark.usefixtures("_dynamic_mock_input")(test_func)
        else:
            # 传入固定列表时，创建静态mock fixture
            @pytest.fixture(autouse=True)
            def _fixed_mock_input(monkeypatch):
                patched_input = _create_patched_input(input_source)
                monkeypatch.setattr("builtins.input", patched_input)
            
            return pytest.mark.usefixtures("_fixed_mock_input")(test_func)
    return decorator

用法示例 #

场景1：直接传入固定字符串列表

def get_name(prompt: str) -> str:
    return input(prompt)

@mock_input(["Alice", "Bob", "Carol"])
def test_get_name_fixed():
    assert get_name("What is your name?") == "Alice"
    assert get_name("What is your name?") == "Bob"
    assert get_name("What is your name?") == "Carol"

场景2：配合@pytest.mark.parametrize使用自定义参数名

@pytest.mark.parametrize(
    "custom_input_list",  # 可使用任意自定义参数名
    (["Alice", "Bob", "Carol"], ["Dale", "Evie", "Frank", "George"]),
)
@mock_input("custom_input_list")  # 传入parametrize定义的参数名
def test_get_name_parametrized(custom_input_list):
    for name in custom_input_list:
        assert get_name("What is your name?") == name

方案优势 #

• 灵活性拉满：配合parametrize时可自由定义参数名，不再受固定名称约束
• 使用简便：只需导入mock_input一个对象，拆分到其他文件时无需额外导入fixture
• 兼容性强：同时支持固定列表和参数化两种测试场景，覆盖绝大多数mock input需求

内容的提问来源于stack exchange，提问作者Ross B.