Python 使用 pytest_mock 在函数中模拟多个查询

Question

我正在为一个包含多个sql查询的函数编写单元测试用例。我正在使用psycopg2模块并尝试模拟光标。

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

测试.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}

但是模拟者总是采用最后一个模拟对象(第二个查询)。我已经尝试了多次黑客攻击，但没有成功。如何在一个函数中模拟多个查询并成功通过单元测试用例？是否可以以这种方式编写单元测试用例，或者我是否需要将查询拆分为不同的函数？

Answer 1

在深入研究文档之后，我能够在 @Pavel Vergeev 建议的 unittest 模拟装饰器和 side_effect 的帮助下实现这一目标。能够编写足以测试功能的单元测试用例。

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}

除此之外，如果查询中有动态参数，也可以通过以下方法断言。

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

因此，当执行第一个查询时，call_args_list[0][0][0]处的cursor.execute(query,(parameter_name,))查询可以获取并断言，在call_args_list[0][0][1]处可以获得第一个参数parameter_name。类似地增加索引，可以获取并断言所有其他参数和不同的查询。