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 ?

Answer 1

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