如何让第三方函数多次返回不同值，同时执行原函数生成输出？

你需要让mock函数前两次返回失败码触发重试，第三次实际执行原函数生成输出，核心是利用pytest-mock的side_effect自定义逻辑，而非单纯返回固定值。

具体实现步骤 #

1. 保存原函数引用：在patch目标函数前，先获取真实的Whitebox工具函数，避免patch后无法调用原逻辑。
2. 自定义side_effect函数：通过计数器记录调用次数，前两次返回失败码1，第三次调用原函数并返回其结果。
3. 替换目标函数：用带自定义side_effect的mock替换原Whitebox函数。

代码示例 #

def test_wbt_retry_logic(mocker):
    # 保存原Whitebox函数的引用
    original_bdlc = whitebox.WhiteboxTools.breach_depressions_least_cost
    call_count = 0

    def mock_side_effect(*args, **kwargs):
        nonlocal call_count
        call_count += 1
        # 前两次调用返回失败码
        if call_count <= 2:
            return 1
        # 第三次调用执行原函数并返回结果
        return original_bdlc(*args, **kwargs)

    # 替换目标函数
    mocker.patch(
        "whitebox.WhiteboxTools.breach_depressions_least_cost",
        side_effect=mock_side_effect
    )

    # 执行你的业务逻辑（调用run_wbt）
    run_wbt(
        whitebox.WhiteboxTools.breach_depressions_least_cost,
        dem_in_path,
        dem_fix_path,
        dist=dist,
        max_cost=max_cost,
        min_dist=min_dist,
        fill=fill,
    )

    # 验证函数确实被调用了3次（触发两次重试）
    assert call_count == 3

通用化封装 #

如果需要对多个Whitebox工具函数做类似测试，可以封装一个通用的mock生成函数，减少重复代码：

def create_retry_mock(original_func, fail_times=2):
    call_count = 0
    def side_effect(*args, **kwargs):
        nonlocal call_count
        call_count += 1
        if call_count <= fail_times:
            return 1
        return original_func(*args, **kwargs)
    return side_effect

使用时只需：

def test_multiple_wbt_retry(mocker):
    original_bdlc = whitebox.WhiteboxTools.breach_depressions_least_cost
    mocker.patch(
        "whitebox.WhiteboxTools.breach_depressions_least_cost",
        side_effect=create_retry_mock(original_bdlc)
    )

    # 执行测试逻辑...

可选调整：每次重试都执行原函数 #

若允许每次重试都执行原函数（覆盖输出也可），可修改side_effect让前两次执行原函数后返回1，第三次返回真实结果：

def mock_side_effect(*args, **kwargs):
    nonlocal call_count
    call_count +=1
    # 每次都执行原函数
    result = original_bdlc(*args, **kwargs)
    if call_count <=2:
        return 1
    return result

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