Python mock multiple queries in a function using pytest_mock

Question

I am writing unit test case for a function which has multiple sql queries in it.I am using psycopg2 module and trying to mock the cursor.

app.py

import psycopg2

def my_function():
    # all connection related code goes here ...

    query = "SELECT name,phone FROM customer WHERE name='shanky'"
    cursor.execute(query)
    columns = [i[0] for i in cursor.description]
    customer_response = []
    for row in cursor.fetchall():
        customer_response.append(dict(zip(columns, row)))

    query = "SELECT name,id FROM product WHERE name='soap'"
    cursor.execute(query)
    columns = [i[0] for i in cursor.description]
    product_response = []
    for row in cursor.fetchall():
        product_response.append(dict(zip(columns, row)))

    return product_response

test.py

from pytest_mock import mocker
import psycopg2

def test_my_function(mocker):
    from my_module import app
    mocker.patch('psycopg2.connect')

    #first query
    mocked_cursor_one = psycopg2.connect.return_value.cursor.return_value
    mocked_cursor_one.description = [['name'],['phone']]
    mocked_cursor_one.fetchall.return_value = [('shanky', '347539593')]
    mocked_cursor_one.execute.call_args == "SELECT name,phone FROM customer WHERE name='shanky'"

    #second query
    mocked_cursor_two = psycopg2.connect.return_value.cursor.return_value
    mocked_cursor_two.description = [['name'],['id']]
    mocked_cursor_two.fetchall.return_value = [('nirma', 12313)]
    mocked_cursor_two.execute.call_args == "SELECT name,id FROM product WHERE name='soap'"

    ret = app.my_function()
    assert ret == {'name' : 'nirma', 'id' : 12313}

But the mocker always takes the last mock object (the second query).I have already tried multiple hacks, but that didn't work out. How can i mock multiple queries in one function and successfully pass the unit test case? Is it possible to write a unit test case in this fashion or do i need to split the queries in different functions?

Answer 1

After drilling a lot through the documentation, I was able to achieve this with the help of unittest mock decorator and side_effect which was suggested by @Pavel Vergeev.I was able to write a unit test case that is good enough to test the functionality.

from unittest import mock
from my_module import app

@mock.patch('psycopg2.connect')
def test_my_function(mocked_db):

    mocked_cursor = mocked_db.return_value.cursor.return_value

    description_mock = mock.PropertyMock()
    type(mocked_cursor).description = description_mock

    fetchall_return_one = [('shanky', '347539593')]

    fetchall_return_two = [('nirma', 12313)]

    descriptions = [
        [['name'],['phone']],
        [['name'],['id']]
    ]

    mocked_cursor.fetchall.side_effect = [fetchall_return_one, fetchall_return_two]

    description_mock.side_effect = descriptions

    ret = app.my_function()

    # assert whether called with mocked side effect objects
    mocked_db.assert_has_calls(mocked_cursor.fetchall.side_effect)

    # assert db query count is 2
    assert mocked_db.return_value.cursor.return_value.execute.call_count == 2

    # first query
    query1 = """
            SELECT name,phone FROM customer WHERE name='shanky'
            """
    assert mocked_db.return_value.cursor.return_value.execute.call_args_list[0][0][0] == query1

    # second query
    query2 = """
            SELECT name,id FROM product WHERE name='soap'
            """
    assert mocked_db.return_value.cursor.return_value.execute.call_args_list[1][0][0] == query2

    # assert the data of response
    assert ret == {'name' : 'nirma', 'id' : 12313}

In addition to this if there are dynamic parameters in the query, that can be asserted too by the following method.

assert mocked_db.return_value.cursor.return_value.execute.call_args_list[0][0][1] = (parameter_name,)

so when the first query is executed, cursor.execute(query,(parameter_name,)) at call_args_list[0][0][0] the query can be obtained and asserted, at call_args_list[0][0][1] the first parameter parameter_name can be obtained. similarly incrementing the index, all the other params and different queries can be obtained and asserted.