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 

what i would like to write some unittests. I am wondering how to use mock.patch to fix cursor and join variables so that they return a fake dataset? I tried the following code segment 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

Since the cursor is the return value of con.cursor , you only need a con.cursor connection, and then configure it correctly. For instance,

 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() 

Answer 2

You have a series of related calls, each of which returns a new object. If you psycopg2.connect() only psycopg2.connect() , you can follow this chain of calls (each of which .return_value dummy objects) through the .return_value attributes that reference the returned .return_value 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) 

Since you hold references to the dummy connect function, as well as dummy connections and cursor objects, you can also claim that 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 do not need to test them, you can simply return_value links to go directly to the result of calling cursor() 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 use the connection as a context manager to automatically commit the transaction and use as to bind the object returned by __enter__() to the new name (for example, with psycopg2.connect(...) as conn: #... ) then you will need to add an extra __enter__.return_value to 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 with conn.cursor() as cursor: conn.cursor.return_value.__enter__.return_value The conn.cursor.return_value.__enter__.return_value is assigned as target object.