pytest - 如何模拟从一个方法中的两个不同游标调用返回一个值答案

【问题标题】：pytest - how to mock return a value from two different cursor calls in a methodpytest - 如何模拟从一个方法中的两个不同游标调用返回一个值
【发布时间】：2023-01-27 17:09:12
【问题描述】：

我必须用 python 编写单元测试。我需要在一个方法中模拟两个不同的游标调用。

sql.py文件

def call_sql(conn, b):
    
    query1 = q1
    query2 = q2

    cur = conn.cursor()
    run1 = cur.execute(query1).fetchone()
    run2 = cur.execute(query2).fetchone()

    count1 = run1[0]
    count2 = run2[0]

    if count1 == count2:
        print('success')
    else:
        print('fail')

def test_Call_sql(self):
    mock_connect = MagicMock()

    connection = mock_connect.return_value
    cursor = connection.cursor.return_value
    cursor.fetchone.return_value = (5,)

问题：如何模拟两个单独的调用？

【问题讨论】：

有没有办法直接模拟 count1 和 count2 ？这样我就可以验证这两个条件？

标签： python mocking pytest magicmock

【解决方案1】：

代码的几个问题。正如@Kun 正确指出的那样，您应该使用side_effect。但是，在模拟时请注意函数调用以及返回的内容和不返回的内容。在您的尝试中，您没有指定模拟游标需要调用 execute 作为其调用之一，并且也有一个返回值。下面是稍微修改后按预期工作的代码。我更改了您的函数以返回一个布尔值，因为这比捕获打印到 stdout 更容易测试。

from unittest.mock import MagicMock
import pytest


def call_sql(conn):
    
    query1 = ""
    query2 = ""

    cur = conn.cursor()
    run1 = cur.execute(query1).fetchone()
    run2 = cur.execute(query2).fetchone()

    count1 = run1[0]
    count2 = run2[0]

    return count1 == count2


@pytest.mark.parametrize(
    "data,expected", 
    [
        ([(1, ), (1, )], True),
        ([(1, ), (2, )], False)
    ]
)
def test_Call_sql(data, expected):
    mock_connect = MagicMock()
    cursor = mock_connect.cursor.return_value
    cursor.execute.return_value.fetchone.side_effect = data

    assert call_sql(mock_connect) is expected

========================================= test session starts ==========================================
platform darwin -- Python 3.8.9, pytest-7.0.1, pluggy-1.0.0
rootdir: **
plugins: mock-3.7.0
collected 2 items                                                                                      

test_script.py ..                                                                                [100%]

===================================== 2 passed in 0.01s =====================================

【讨论】：

谢谢，它起作用了，我在执行时缺少 return_value

【解决方案2】：

您可以编写一个函数来根据 MagicMock 内部副作用的参数确定输出

def side_effect_func(val):
    if val == query1:
        # do sth
    if val == query2:
        # do sth

m = MagicMock(side_effect=side_effect_func)

【讨论】：

【解决方案3】：

你可以使用side effect

在您的示例中，您将替换

cursor.fetchone.return_value = (5,0)

和

cursor.fetchone.side_effect = [(5,0), (6,0)]  # 1st call returns (5,0), 2nd call returns (6,0)

【讨论】：

试过了，但它仍然为 count1 和 count2 提供了一些随机模拟 ID，它没有根据需要分配值 5 或 6
请帮助解决我的问题