У меня есть этот сегмент кода в 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'
Поскольку курсор является возвращаемым значением 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()
У вас есть ряд связанных вызовов, каждый из которых возвращает новый объект. Если вы 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 целевой объект.