Mocking Only One Function of a Class

Question

I would like to write a unit test that can ensure a SQL statement in a function call is schematically correct. It should test the execution of this call. I would then like to mock the call to commit, so that no insertions to the database take place. I'm using psycopg2 for my tests.

I have a function like:

def test_insert(a, b, c):
    con = psycopg2.connect(os.environ['PGDB'])
    cur = con.cursor()
    cur.execute('insert into test_table values ({a}, {b}, {c})'.format(a=a, b=b, c=c))
    con.commit()
    con.close()

when calling test_insert(1,2,3) I see the row inserted into the table. Now I try to mock the call. I've taken a few approaches so far:

@mock.patch('psycopg2.connect')
def test(mock_connect, a, b, c):
    mock_con = mock_connect.return_value
    mock_con.commit.return_value = None
    insert_row(a, b, c)

This seems to work but does not actually call the execution statement. test_insert(1,4,'xyz') fails for instance while test(1,4,'xyz') does not. Next I tried to mock just the commit method of the connection class in psycopg2:

@mock.patch('psycopg2.extensions.connection.commit')
def test_insert(mock_commit, a, b, c):
    mock_commit.return_value = None
    insert_row(a,b,c)

but this gives me a syntax error

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/a/.virtualenv/test/lib/python2.7/site-packages/mock/mock.py", line 1318, in patched
    patching.__exit__(*exc_info)
  File "/home/a/.virtualenv/test/lib/python2.7/site-packages/mock/mock.py", line 1480, in __exit__
    setattr(self.target, self.attribute, self.temp_original)
TypeError: can't set attributes of built-in/extension type 'psycopg2.extensions.connection'

Is there a good way to do what I am trying to do?

Answer 1

I assume you are using pytest and it is not a good practice to name your functions starting with test_ if they are not actual tests as this will probably raise problems with the testing framework. Therefore I slightly modified your initial snippet as follows and I named the module psyco.py

import psycopg2
import os

def insert(a, b, c):
    con = psycopg2.connect(os.environ['PGDB'])
    import ipdb; ipdb.set_trace()
    cur = con.cursor()
    cur.execute('insert into test_table values ({a}, {b}, {c})'.format(a=a, b=b, c=c))
    con.commit()
    con.close()

Next, I created the test for your method by taking into account how patch works and where to patch. As you are dealing with os environment variables this question can help you understand why I mocked it that way.

An example implementation of the test could be as follows:

from psyco import insert
from unittest.mock import patch, Mock, MagicMock
import os

@patch.dict(os.environ,{'PGDB':'db_url'})
@patch('psycopg2.connect')
def test_insert_function(psycopg2_mock):

    x = 1
    y = 4
    z = 'xyz'
    sql_query = 'insert into test_table values ({0}, {1}, {2})'.format(x,y,z)
    insert(x,y,z)
    assert psycopg2_mock.return_value.cursor.call_count == 1
    psycopg2_mock.return_value.cursor.return_value.execute.assert_called_with(sql_query)
    assert psycopg2_mock.return_value.commit.call_count == 1
    assert psycopg2_mock.return_value.close.call_count == 1