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

Question

我必须用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,)

问题:如何模拟两个独立的调用?

Answer 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 =====================================

Answer 2

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

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

Answer 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)