Python单元测试如何mock函数实现根据入参返回对应值

完全可以实现，你可以通过unittest.mock的side_effect属性实现根据入参返回指定值的效果，这也是处理多参数不同返回值Mock场景的常用方案。

正确实现示例 #

你示例中直接给return_value传字典的写法是错误的，会导致每次调用get_name_by_id都返回整个字典，而非对应ID的名称。参考以下两种写法修改测试代码即可：

写法1：上下文管理器形式

# file:[test/test_my_module.py]
from unittest import mock
from my_module import order_names

def test_order_names():
    # 提前定义ID到名称的映射关系
    id_name_map = {"3": "Mark", "4": "Kate", "5":"Alfred"}
    # patch替换目标函数，将side_effect绑定为字典的get方法
    with mock.patch("src.my_module.get_name_by_id", side_effect=id_name_map.get) as mock_get_name:
        ordered_names = order_names(["3", "4", "Suzan", "5"])
        # 可选：额外断言mock被正确调用，参数符合预期
        mock_get_name.assert_has_calls(
            [mock.call("3"), mock.call("4"), mock.call("5")],
            any_order=True
        )
    assert ordered_names == "Alfred, Kate, Mark, Suzan"

写法2：装饰器形式

# file:[test/test_my_module.py]
from unittest import mock
from my_module import order_names

@mock.patch("src.my_module.get_name_by_id")
def test_order_names(mock_get_name):
    id_name_map = {"3": "Mark", "4": "Kate", "5":"Alfred"}
    # 给mock对象手动指定side_effect
    mock_get_name.side_effect = id_name_map.get
    
    ordered_names = order_names(["3", "4", "Suzan", "5"])
    assert ordered_names == "Alfred, Kate, Mark, Suzan"

原理解释 #

• side_effect传入可调用对象时，每次Mock被调用，都会将调用时的入参传递给这个可调用对象，返回它的执行结果。上面的写法用字典的get方法，正好匹配get_name_by_id接收一个id参数的调用形式，自动根据入参返回对应名称。
• 如果需要处理更复杂的返回逻辑（比如无效ID抛异常、参数校验等），可以自定义函数作为side_effect，示例：

def mock_get_name_logic(id):
    id_name_map = {"3": "Mark", "4": "Kate", "5":"Alfred"}
    if id not in id_name_map:
        raise ValueError("无效ID")
    return id_name_map[id]

# 直接赋值给side_effect即可
mock_get_name.side_effect = mock_get_name_logic

注意事项 #

• 确保patch的路径正确：你需要Patch的是order_names函数实际调用的get_name_by_id的路径，本例中两个函数在同一个模块，写src.my_module.get_name_by_id是正确的。
• 如果需要给不存在的ID设置默认返回值，可以调整side_effect为lambda表达式：side_effect=lambda x: id_name_map.get(x, "未知用户")。

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