如何让Mock的BigQueryConnector.query_to_pandas两次返回不同DataFrame

方法1：按调用顺序返回不同DataFrame（简单场景） #

如果两次调用不需要区分传入的query，只需要按调用顺序返回不同结果，直接利用mock的side_effect参数传一个包含两个DataFrame的列表即可，第一次调用返回列表第一个元素，第二次返回第二个。示例代码：

import pandas as pd
from unittest.mock import patch
from your_module import your_function  # 替换为你的被测试函数所在模块

def test_your_function():
    # 构造两个不同的Mock DataFrame
    df1 = pd.DataFrame({"col1": [1, 2], "col2": ["a", "b"]})
    df2 = pd.DataFrame({"col3": [3, 4], "col4": ["c", "d"]})

    # 注意：patch路径必须指向被测试函数中实际引用的BigQueryConnector，而非原类定义路径
    @patch("your_module.BigQueryConnector.query_to_pandas")
    def inner_test(mock_query):
        # 设置side_effect，按调用顺序返回指定DataFrame
        mock_query.side_effect = [df1, df2]
        
        # 执行被测试函数
        result = your_function()
        
        # 验证调用次数和结果（根据你的业务逻辑补充断言）
        assert mock_query.call_count == 2

    inner_test()

方法2：根据传入的query返回对应DataFrame（复杂场景） #

如果两次调用的query内容不同，需要根据query返回对应DataFrame，可将side_effect设置为自定义函数，通过判断传入的query参数返回不同结果：

import pandas as pd
from unittest.mock import patch
from your_module import your_function

def test_your_function_with_different_queries():
    # 为不同query构造对应Mock DataFrame
    df_user_data = pd.DataFrame({"user_id": [101, 102], "name": ["Alice", "Bob"]})
    df_order_data = pd.DataFrame({"order_id": [201, 202], "amount": [100, 200]})

    def mock_query_handler(query):
        # 根据实际query内容匹配返回结果
        if "SELECT * FROM user_table" in query:
            return df_user_data
        elif "SELECT * FROM order_table" in query:
            return df_order_data
        else:
            # 可选：处理未匹配的query，比如返回空DataFrame或抛出异常
            return pd.DataFrame()

    @patch("your_module.BigQueryConnector.query_to_pandas")
    def inner_test(mock_query):
        mock_query.side_effect = mock_query_handler
        
        result = your_function()
        
        # 验证两次调用的query参数是否符合预期
        mock_query.assert_any_call("SELECT * FROM user_table")
        mock_query.assert_any_call("SELECT * FROM order_table")
        # 补充其他业务逻辑断言

    inner_test()

关键注意事项 #

• Patch路径正确性：必须mock被测试函数所在模块中导入的BigQueryConnector，而非BigQueryConnector原定义的模块路径。例如，若被测试函数在my_app/data_process.py中通过from connectors.bigquery import BigQueryConnector导入类，那么patch路径应为"my_app.data_process.BigQueryConnector.query_to_pandas"，否则mock不会生效。
• side_effect扩展用法：除了列表和自定义函数，还可以传入异常类/实例，用来模拟调用失败的场景。

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