Python单元测试：同一MySQLClient对象多函数返回不同fetchall值的实现

问题描述 #

在集成测试中已Mock MySQLClient对象，首次调用fetchall可返回预设值，但需要在不同函数（func1/func2/func3）中调用同一对象的fetchall时返回不同结果，现有代码结构如下：

测试函数代码 #

@mock.patch("src.calculate.get_mysql_connection", return_value = mock.Mock())
@mock.patch("src.calculate.get_mongo_connection", return_value = mock.Mock())
def test_initiate_calc(self, mongo_patcher, mysql_patcher):
    mysql_patcher.return_value.\
            execute_query.return_value.\
            fetchall.return_value = [1,2,3,4]
    
    test_main()

主函数代码 #

def main():
    mongo_client = get_mongo_connection() 
    mysql_client = get_mysql_connection()  # mocked mysql object
    func1(mysql_client)
    func2(mysql_client)
    func3(mysql_client)
    func4(mysql_client)

def func1(mysql_client):
    query = "select something"
    params = (some_params, )
    cursor = mysql_client.execute_query(query, params)
    result1 = cursor.fetchall()  # 当前返回[1,2,3,4]

def func2(mysql_client):
    query = "select something else"
    params = (some_params, )
    cursor = mysql_client.execute_query(query, params)
    result2 = cursor.fetchall()  # 需要返回另一组值

def func3(mysql_client):
    query = "select another thing"
    params = (some_params, )
    cursor = mysql_client.execute_query(query, params)
    result3 = cursor.fetchall()  # 需要返回第三组值

解决方案 #

有两种常用方式实现需求，可根据测试场景选择：

方法1：按调用顺序返回不同值 #

如果无需区分查询语句，仅按fetchall的调用顺序返回结果，可使用side_effect属性，传入结果列表后每次调用会依次取列表中的元素：

@mock.patch("src.calculate.get_mysql_connection", return_value = mock.Mock())
@mock.patch("src.calculate.get_mongo_connection", return_value = mock.Mock())
def test_initiate_calc(self, mongo_patcher, mysql_patcher):
    # 按func1、func2、func3的调用顺序定义返回值
    fetchall_results = [
        [1,2,3,4],   # func1的返回值
        [5,6,7,8],   # func2的返回值
        [9,10,11,12] # func3的返回值
    ]
    # 给fetchall设置side_effect实现顺序返回
    mysql_patcher.return_value.execute_query.return_value.fetchall.side_effect = fetchall_results
    
    test_main()

方法2：根据查询语句/参数返回不同值 #

如果需要根据不同的查询语句或参数返回对应结果，可给execute_query设置自定义side_effect，根据传入参数返回不同的cursor mock：

@mock.patch("src.calculate.get_mysql_connection", return_value = mock.Mock())
@mock.patch("src.calculate.get_mongo_connection", return_value = mock.Mock())
def test_initiate_calc(self, mongo_patcher, mysql_patcher):
    def mock_execute_query(query, params):
        cursor = mock.Mock()
        # 根据查询语句匹配对应返回值
        if query == "select something":
            cursor.fetchall.return_value = [1,2,3,4]
        elif query == "select something else":
            cursor.fetchall.return_value = [5,6,7,8]
        elif query == "select another thing":
            cursor.fetchall.return_value = [9,10,11,12]
        else:
            cursor.fetchall.return_value = []  # 默认返回值
        return cursor
    
    # 给execute_query绑定自定义逻辑
    mysql_patcher.return_value.execute_query.side_effect = mock_execute_query
    
    test_main()

补充说明 #

• side_effect既可以接收列表，也可以接收自定义函数，用于处理更复杂的匹配逻辑
• 如果func4也需要特定返回值，只需在列表或条件判断中添加对应结果即可

内容的提问来源于stack exchange，提问作者Just another object