How to mock psycopg2 cursor object?
Solution 1
Since the cursor is the return value of con.cursor
, you only need to mock the connection, then configure it properly. For example,
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()
Solution 2
You have a series of chained calls, each returning a new object. If you mock just the psycopg2.connect()
call, you can follow that chain of calls (each producing mock objects) via .return_value
attributes, which reference the returned mock for such calls:
@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)
Because you hold onto references for the mock connect
function, as well as the mock connection and cursor objects you can then also assert if they were called correctly:
mock_connect.assert_called_with(**connection_stuff)
mock_con.cursor.asset_called_with(cursor_factory=DictCursor)
mock_cur.execute.assert_called_with("Super duper SQL query")
If you don't need to test these, you could just chain up the return_value
references to go straight to the result of cursor()
call on the connection object:
@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)
Note that if you are using the connection as a context manager to automatically commit the transaction and you use as
to bind the object returned by __enter__()
to a new name (so with psycopg2.connect(...) as conn: # ...
) then you'll need to inject an additional __enter__.return_value
in the call chain:
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()
The same applies to the result of with conn.cursor() as cursor:
, the conn.cursor.return_value.__enter__.return_value
object is assigned to the as
target.
Solution 3
The following answer is the variation of above answers.
I was using django.db.connections
cursor object.
So following code worked for me
@patch('django.db.connections')
def test_supercool_method(self, mock_connections):
query_result = [("field1a", "field2a"), ("field1b", "field2b")]
mock_connections.__getitem__.return_value.cursor.return_value.__enter__.return_value.fetchall.return_value = query_result
result = supercool_method()
self.assertIsInstance(result, list)
Related videos on Youtube
zyshara
Updated on July 11, 2022Comments
-
zyshara almost 2 years
I have this code segment in 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
that I'd like to write some unittests for. I'm wondering how to use
mock.patch
in order to patch out the cursor and connection variables so that they return a fake set of data? I've tried the following segment of code for my unittests but to no avail:@mock.patch("psycopg2.connect") @mock.patch("psycopg2.extensions.cursor.fetchall") def test_super_awesome_stuff(self, a, b): testing = super_cool_method()
But I seem to get the following error:
TypeError: can't set attributes of built-in/extension type 'psycopg2.extensions.cursor'
-
allenlin1992 almost 8 yearsThis doesn't work for me. The result I get from fetchall is '<MagicMock name='connect().cursor().fetchall()' id='40430416'>'
-
chepner almost 8 years@AllenLin That's because I didn't configure a return value for
fetchall
, justexecute
. -
allenlin1992 almost 8 yearscursor.execute usually returns None. Why would you configure a return value for execute but not fetchall?
-
chepner almost 8 yearsBecause clearly I should have mocked
fetchall
:) Sorry about the confusion. -
allenlin1992 almost 8 yearsUnfortunately still not working for me. I've posted my problem as it's own question: stackoverflow.com/questions/37164891/…
-
Martijn Pieters almost 8 years@chepner: actually, because both
connect
andexecute
are called, you need to inject a few morereturn_value
references here. Not that the return value ofexecute
is used though,fetchall
is never referenced on that mock. -
Davos over 5 yearsThis is awesome, thanks for the context manager
__enter__
advice. From a philosophy perspective, is this a suggested way of testing? it seems like it is just comparingexpected
toexpected
as returned from the mock, so this test would always work even if the logic insuper_cool_method()
changed - as long as the syntax is valid the test would never break. Is this a simple example to illustrate the point, or is it doing something I've missed? I can see that you can enforceDictCursor
which is valuable, but can't see whyfetchall.return_value
is being tested, trying to learn. -
Martijn Pieters over 5 years@Davos yes, this is a good philosophy because your unit test should not be testing if psycopg2 is working correctly, it should be testing if your code is using the results correctly. Real-world code under test will do something with the result from
fetchall()
and it is that something that matters. Or you want to test what was passed to cursor.execute(), etc. -
Andrey Semakin about 5 yearsCorrect line is
mock_connect.return_value.cursor.return_value.execute.return_value.fetch_all.return_value = query_result
-
DarioB almost 5 yearsFor future reference, this doesn't work because he mistyped
fetch_all
, when he should have writefetchall
. Changing to that it woks. -
BHC almost 5 yearsCorrect line is actually
mock_connect.return_value.cursor.return_value.fetch_all.return_value = query_result
-fetchall
is not a member ofexecute
-
chepner almost 5 yearsI think it's fixed now.
fetchall
is, indeed, not a method of the return value ofexecute
; ratherexecute
prepares a cursor for usingfetchall
. -
Martijn Pieters almost 2 years@RichardA: Thanks for the edit! I thought that perhaps back in 2016 and Python 2.7, the (separately installed)
mock
project hadcalled_with
as an alias forassert_called_with
, but alas, it was simply me getting it wrong. -
Richard A almost 2 years@MartijnPieters: and thank you for your answer from way back then. You saved me a lot of time and fumbling around.