如何用Pytest Mock返回含render方法的对象以测试代码？

问题背景 #

原业务代码中，get_query是自定义函数，调用后会返回一个带有render方法的对象，代码如下：

query = get_query(f"export_user_{user_component}").render()

需要在测试中Mock这个get_query函数，但遇到了问题。

现有测试代码 #

测试代码使用mock.patch装饰器Mock相关函数：

@mock.patch("de_etl.operations.user_data_export.components.get_query", side_effect=mock_get_query)
@mock.patch("de_etl.operations.user_data_export.util.execute_query", side_effect=mock_execute_query)
def test__export_user_empty_profile(self, execute_query, get_query):
    with pytest.raises(Exception) as exc:
        de_etl.operations.user_data_export.components._export_user(user_id="12345678910", db=mock.Mock())

    assert exc.value.args[0] == "User Profile must exist"

错误原因 #

最初的mock_get_query直接返回字符串，导致调用.render()时触发'str' object has no attribute 'render'错误：

def mock_get_query(name):
    return name

需要让mock_get_query返回一个带有render方法的对象。

自行尝试的方案 #

编写了如下代码，询问是否可行：

def mock_get_query(sql_to_mock):
    def mock_get_query_called(actual_sql):
        if sql_to_mock == actual_sql:
            mocked_query = mock.Mock()
            mocked_query.render.return_value = sql_to_mock
            return mocked_query
        else:
            mocked_query = mock.Mock()
            mocked_query.render.return_value = actual_sql
            return mocked_query

    return mock_get_query_called

你写的这个方案是可行的，但可以简化实现，让代码更简洁直观：

简化实现1：直接返回带render方法的Mock对象 #

不需要嵌套函数，直接在mock_get_query中创建Mock对象并设置render方法的返回值：

def mock_get_query(name):
    mock_obj = mock.Mock()
    mock_obj.render.return_value = name
    return mock_obj

无论传入什么name，返回的Mock对象调用render()都会返回对应的name，完全满足原业务代码的调用需求。

简化实现2：针对特定查询做自定义返回 #

如果需要对特定的查询字符串返回不同结果，可以直接在函数中判断：

def mock_get_query(name):
    mock_obj = mock.Mock()
    # 针对特定查询返回自定义SQL
    if name == "export_user_target_component":
        mock_obj.render.return_value = "SELECT * FROM target_user_table"
    else:
        mock_obj.render.return_value = name
    return mock_obj

测试代码调整 #

保持原有的@mock.patch装饰器不变，替换mock_get_query为上述简化版本即可，测试函数的其他部分无需修改。

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