pytest – how to mock return a value from two different cursor calls in a method

Question

I have to write unit test in python. I need to mock two different cursor calls in a single method.

sql.py file

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

Question: how to mock two separate calls ?

Asked By: viki

||

Source

Answer 1

You can write a function to determine the output based on arguments on the side effect inside the MagicMock

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

m = MagicMock(side_effect=side_effect_func)

Answered By: seel

Answer 2

You can use side effect

In your example, you will replace

cursor.fetchone.return_value = (5,0)

with

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

Answered By: Kun

Answer 3

Several issues with the code. As @Kun correctly points out, you should use side_effect. However, when mocking take note of the function calls and what returns and what doesn’t. In your attempts you do not specify that the mock cursor needs to call execute as one of its calls, and have a return value for that as well. Below, is the code slightly modified which works as expected. I changed your function to return a boolean as that is easier to test than capturing printing to 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 =====================================

Answered By: gold_cy

pytest – how to mock return a value from two different cursor calls in a method

Question:

Answers: