pytest如何mock单个方法内两次不同cursor调用的返回值

你可以通过side_effect属性实现多次调用返回不同结果的需求，有两种常用实现方案：

方案1：按调用顺序返回固定值（写法最简单） #

MagicMock的side_effect传入可迭代对象时，会在每次调用mock方法时按顺序返回迭代器中的下一个值，刚好匹配代码中两次fetchone()顺序执行的逻辑：

from unittest.mock import MagicMock
import unittest
from sql import call_sql

class TestSqlCase(unittest.TestCase):
    def test_Call_sql(self):
        mock_connect = MagicMock()
        connection = mock_connect.return_value
        cursor = connection.cursor.return_value
        # 按fetchone的调用顺序配置返回值：第一次调用返回(5,)，第二次返回(10,)
        # 按需修改两个元组的值即可覆盖相等/不等的逻辑分支
        cursor.fetchone.side_effect = [(5,), (10,)]

        # 传入mock连接调用被测方法
        call_sql(connection, b="test_val")

        # 可选：添加执行断言验证逻辑
        self.assertEqual(cursor.execute.call_count, 2)

方案2：根据传入的SQL语句匹配返回值（鲁棒性更强） #

如果不想依赖代码中的SQL执行顺序，想根据execute()接收到的SQL参数动态返回对应结果，可以给execute方法配置自定义side_effect逻辑：

from unittest.mock import MagicMock
import unittest
from sql import call_sql

class TestSqlCase(unittest.TestCase):
    def test_Call_sql(self):
        mock_connect = MagicMock()
        connection = mock_connect.return_value
        cursor = connection.cursor.return_value

        def mock_execute_logic(query):
            # 根据传入的SQL语句返回配置好对应fetchone结果的mock对象
            exec_result_mock = MagicMock()
            if query == q1: # 替换为query1对应的实际SQL值
                exec_result_mock.fetchone.return_value = (5,)
            elif query == q2: # 替换为query2对应的实际SQL值
                exec_result_mock.fetchone.return_value = (10,)
            return exec_result_mock
        
        cursor.execute.side_effect = mock_execute_logic

        # 调用被测方法
        call_sql(connection, b="test_val")

        # 可选：断言两条SQL都被执行
        cursor.execute.assert_any_call(q1)
        cursor.execute.assert_any_call(q2)

注意事项 #

• 方案1适合SQL执行顺序固定、后续不会调整的场景，代码量最少；方案2不受执行顺序影响，就算后续调整两条SQL的调用顺序，mock逻辑依然有效，更适合复杂业务场景。
• 需要测试count1 == count2分支时，只需要把两个返回的元组值设为相同即可，比如都配置为(5,)，就会触发print('success')逻辑。

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