如何在Pytest中为循环内的同一函数Mock不同返回值？

现有处理汽车喷漆的函数：

def do_operate(car_list: list):
    flawless_car_list = list()
    for car in car_list:
        if paint_the_car(car):
            flawless_car_list.append(car)
        else:
            print_log(f"{car} couldn't get painted")

    return flawless_car_list

paint_the_car可能随机返回False，需要测试第二辆车喷漆失败，其余成功的场景，确保程序继续处理后续车辆并记录日志。快乐路径测试已实现：

def test_happy_path(mocker):
    test_list = ["honda", "bmw", "ford"]
    mocker.patch("pytest1.paint_the_car", return_value=True)
    mocker.patch("pytest1.print_log")

    card_list = do_operate(car_list=test_list)

    assert len(card_list) == 3
    paint_the_car.assert_called()
    print_log.assert_not_called()

待修改的测试函数需要满足断言：

def test_error_occured_on_second_item(mocker):
    test_list = ["honda", "bmw", "ford"]

    # HOW TO CHANGE MOCKS?
    mocker.patch("pytest1.paint_the_car", return_value=True)
    mocker.patch("pytest1.print_log")

    card_list = do_operate(car_list=test_list)

    assert len(card_list) == 2
    paint_the_car.assert_called()
    print_log.assert_called_with("bmw couldn't get painted")

解决方案 #

方法1：使用side_effect按顺序返回值 #

mock.patch的side_effect参数可以接受一个列表，Mock会按调用顺序依次返回列表中的值。针对这个场景，传入[True, False, True]即可对应三辆车的返回结果：

def test_error_occured_on_second_item(mocker):
    test_list = ["honda", "bmw", "ford"]

    # 修改paint_the_car的Mock，按顺序返回True、False、True
    mocker.patch("pytest1.paint_the_car", side_effect=[True, False, True])
    mock_print_log = mocker.patch("pytest1.print_log")

    card_list = do_operate(car_list=test_list)

    assert len(card_list) == 2
    # 验证paint_the_car被按顺序调用三次
    paint_the_car.assert_has_calls([
        mocker.call("honda"),
        mocker.call("bmw"),
        mocker.call("ford")
    ])
    mock_print_log.assert_called_with("bmw couldn't get painted")

方法2：使用自定义函数根据参数返回 #

如果需要更灵活的逻辑（比如根据车辆名称判断返回值），可以给side_effect传入自定义函数：

def test_error_occured_on_second_item(mocker):
    test_list = ["honda", "bmw", "ford"]

    # 自定义函数：仅当车辆是"bmw"时返回False，其余返回True
    def mock_paint(car):
        return car != "bmw"

    mocker.patch("pytest1.paint_the_car", side_effect=mock_paint)
    mock_print_log = mocker.patch("pytest1.print_log")

    card_list = do_operate(car_list=test_list)

    assert len(card_list) == 2
    mock_print_log.assert_called_with("bmw couldn't get painted")

关键说明 #

• side_effect比return_value更适合动态返回值场景，支持列表、函数、异常等多种类型。
• 方法1适合固定顺序的返回需求，方法2适合需要根据输入参数判断返回值的场景，按需选择即可。

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