Как издеваться над объектом курсора psycopg2?

У меня есть этот сегмент кода в Python2:

def super_cool_method():
    con = psycopg2.connect(**connection_stuff)
    cur = con.cursor(cursor_factory=DictCursor)
    cur.execute("Super duper SQL query")
    rows = cur.fetchall()

    for row in rows:
        # do some data manipulation on row
    return rows

что я хотел бы написать некоторые unittests. Мне интересно, как использовать mock.patch для исправления переменных курсора и соединения, чтобы они возвращали поддельный набор данных? Я пробовал следующий сегмент кода для своих unittests, но безрезультатно:

@mock.patch("psycopg2.connect")
@mock.patch("psycopg2.extensions.cursor.fetchall")
def test_super_awesome_stuff(self, a, b):
    testing = super_cool_method()

Но я, кажется, получаю следующую ошибку:

TypeError: can't set attributes of built-in/extension type 'psycopg2.extensions.cursor'

Ответ 1

Поскольку курсор является возвращаемым значением con.cursor, вам нужно только con.cursor соединение, а затем правильно его настроить. Например,

query_result = [("field1a", "field2a"), ("field1b", "field2b")]
with mock.patch('psycopg2.connect') as mock_connect:
    mock_connect.cursor.return_value.fetchall.return_value = query_result
    super_cool_method()

Ответ 2

У вас есть ряд связанных вызовов, каждый из которых возвращает новый объект. Если вы psycopg2.connect() только psycopg2.connect(), вы можете следовать этой цепочке вызовов (каждый из которых .return_value фиктивные объекты) через атрибуты .return_value, которые ссылаются на возвращенный .return_value для таких вызовов:

@mock.patch("psycopg2.connect")
def test_super_awesome_stuff(self, mock_connect):
    expected = [['fake', 'row', 1], ['fake', 'row', 2]]

    mock_con = mock_connect.return_value  # result of psycopg2.connect(**connection_stuff)
    mock_cur = mock_con.cursor.return_value  # result of con.cursor(cursor_factory=DictCursor)
    mock_cur.fetchall.return_value = expected  # return this when calling cur.fetchall()

    result = super_cool_method()
    self.assertEqual(result, expected)

Поскольку вы держите ссылки на функцию фиктивного connect, а также на фиктивные соединения и объекты курсора, вы также можете утверждать, что они были вызваны правильно:

mock_connect.assert_called_with(**connection_stuff)
mock_con.cursor.called_with(cursor_factory=DictCursor)
mock_cur.execute.called_with("Super duper SQL query")

Если вам не нужно тестировать их, вы можете просто return_value ссылки return_value чтобы перейти непосредственно к результату вызова cursor() объекта соединения:

@mock.patch("psycopg2.connect")
def test_super_awesome_stuff(self, mock_connect):
    expected = [['fake', 'row', 1], ['fake', 'row' 2]]
    mock_connect.return_value.cursor.return_value.fetchall.return_value = expected

    result = super_cool_method()
    self.assertEqual(result, expected)

Обратите внимание, что если вы используете соединение в качестве менеджера контекста для автоматической фиксации транзакции и используете as для привязки объекта, возвращаемого __enter__() к новому имени (например, with psycopg2.connect(...) as conn: #...) тогда вам нужно будет добавить дополнительное __enter__.return_value в цепочку вызовов:

mock_con_cm = mock_connect.return_value  # result of psycopg2.connect(**connection_stuff)
mock_con = mock_con_cm.__enter__.return_value  # object assigned to con in with ... as con    
mock_cur = mock_con.cursor.return_value  # result of con.cursor(cursor_factory=DictCursor)
mock_cur.fetchall.return_value = expected  # return this when calling cur.fetchall()

То же самое относится и к результату with conn.cursor() as cursor: conn.cursor.return_value.__enter__.return_value Объект conn.cursor.return_value.__enter__.return_value назначается as целевой объект.