CODEWITHSUNDEEP
pythonmockingpytestmagicmock
viki
viki

Reputation: 157

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

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 ?

Upvotes: 0

Views: 1797

Answers (3)

gold_cy
gold_cy

Reputation: 14236

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

Upvotes: 2

Kun
Kun

Reputation: 99

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)

Upvotes: 0

seel
seel

Reputation: 135

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)

Upvotes: 0

Related Questions