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'

score 65 · Answer 1 · edited May 30 '22 at 00:35

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.

This 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 comparing `expected` to `expected` as returned from the mock, so this test would always work even if the logic in `super_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 enforce `DictCursor` which is valuable, but can't see why `fetchall.return_value` is being tested, trying to learn. — Davos, Nov 15 '18 at 05:30
@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. — Martijn Pieters, Nov 15 '18 at 08:17

chepner · Accepted Answer · 2019-06-11T14:06:38.533

17

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

edited Jun 11 '19 at 14:06

answered Feb 02 '16 at 02:37

chepner

497,756
71
530
681

5

This doesn't work for me. The result I get from fetchall is '' – allenlin1992 May 11 '16 at 13:29
@AllenLin That's because I didn't configure a return value for `fetchall`, just `execute`. – chepner May 11 '16 at 13:32
1

cursor.execute usually returns None. Why would you configure a return value for execute but not fetchall? – allenlin1992 May 11 '16 at 13:35
1

Because clearly I *should* have mocked `fetchall` :) Sorry about the confusion. – chepner May 11 '16 at 13:36
1

Unfortunately still not working for me. I've posted my problem as it's own question: http://stackoverflow.com/questions/37164891/cant-mock-cursor-fetchone-return-value-python – allenlin1992 May 11 '16 at 13:57
1

@chepner: actually, because both `connect` and `execute` are **called**, you need to inject a few more `return_value` references here. Not that the return value of `execute` is used though, `fetchall` is never referenced on that mock. – Martijn Pieters May 13 '16 at 08:33
Correct line is `mock_connect.return_value.cursor.return_value.execute.return_value.fetch_all.return_value = query_result` – Andrey Semakin Feb 16 '19 at 15:47
For future reference, this doesn't work because he mistyped `fetch_all`, when he should have write `fetchall`. Changing to that it woks. – DarioB May 21 '19 at 09:50
Correct line is actually `mock_connect.return_value.cursor.return_value.fetch_all.return_value = query_result` - `fetchall` is not a member of `execute` – BHC Jun 11 '19 at 14:04
I think it's fixed now. `fetchall` is, indeed, not a method of the return value of `execute`; rather `execute` prepares a cursor for using `fetchall`. – chepner Jun 11 '19 at 14:07

score 2 · Answer 3 · answered Dec 01 '22 at 11:34

2

@patch("psycopg2.connect")
async def test_update_task_after_launch(fake_connection):
    """
    
    """
    fake_update_count =4
    fake_connection.return_value = Mock(cursor=lambda : Mock(execute=lambda x,y :"",  
fetch_all=lambda:['some','fake','rows'],rowcount=fake_update_count,close=lambda:""))

answered Dec 01 '22 at 11:34

Otobong Jerome

401
6
5

more detailed answer: connection_mock.return_value = Mock( __enter__=lambda dsn: Mock( cursor=lambda: Mock( __enter__=lambda cursor_factory: Mock( __enter__=lambda: "", execute=lambda query_str, args: "", fetchone=lambda: [0], ), __exit__=lambda arg1, arg2, arg3, arg4: "", ), ), __exit__=lambda arg1, arg2, arg3, arg4: "", ) – Soyuzbek Orozbek Uulu May 23 '23 at 11:01

score 1 · Answer 4 · answered Apr 20 '20 at 09:08

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)

How to mock psycopg2 cursor object?

4 Answers4

Linked

Related