How to mock psycopg2 cursor object?

Question

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' 

Answer 1

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.called_with(cursor_factory=DictCursor) mock_cur.execute.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.