Python Mocking db connection/unknown type in unit test

Question

Newby to python here. My class uses a database connection to wrap some functions. I have figured out some basic examples successfully. For the more complex library that I am working with, I cannot find close examples of mocking the database connection. In mine, the

class DBSAccess():
    def __init__(self, db_con):
        self.db_con = db_con

    def get_db_perm(self, target_user):
      ## this is where I start having trouble
        with self.db_con.cursor() as cursor:
            cursor.execute("SELECT CAST(sum(maxperm) AS bigint) \
            FROM dbc.diskspace \
            WHERE  databasename = '%s' \
            GROUP BY databasename" % (target_user))

            res = cursor.fetchone()
            if res is not None:
                return res[0]
            else:
                msg = target_user + " does not exist"
                return msg

where db_con is a teradata.UdaExec returns a connection

udaExec = teradata.UdaExec (appName="whatever", version="1.0", logConsole=True)
db_con = udaExec.connect(method="odbc", system='my_sys', username='my_name', password='my_pswd')
dbc_instance = tdtestpy.DBSaccess (db_con)

So for my test to not use any real connection, I have to mock some things out. I tried this combination:

class DBAccessTest(unittest.TestCase):
  def test_get_db_free_perm_expects_500(self):
    uda_exec = mock.Mock(spec=teradata.UdaExec)
    db_con = MagicMock(return_value=None)
    db_con.cursor.fetchone.return_value = [500]
    uda_exec.connect.return_value = db_con
    self.dbc_instance = DBSAccess(db_con)
    self.assertEqual(self.dbc_instance.get_db_free_perm("dbc"), 500)

but my result is messed up because fetchone is returning a mock, not the [500] one item list I was expecting:

AssertionError: <MagicMock name='mock.connect().cursor().[54 chars]312'> != 500

I've found some examples where there is a 'with block' for testing an OS operation, but nothing with database. Plus, I don't know what data type the db_con.cursor is so I can't spec that precisely - I think the cursor is found in UdaExecConnection.cursor() found at Teradata/PyTd.

I need to know how to mock the response that will allow me to test the logic within my method.

Answer 1

The source of your problem is in the following line:

with self.db_con.cursor() as cursor:

with lines calls __enter__ method, which generate in your case a new mock.

The solution is to mock __enter__ method:

db_con.cursor.return_value.__enter__.return_value = cursor

Your tests:

class DBAccessTest(unittest.TestCase):
    def test_get_db_free_perm_expects_500(self):
        db_con = MagicMock(UdaExecConnection)
        cursor = MagicMock(UdaExecCursor)
        cursor.fetchone.return_value = [500]
        db_con.cursor.return_value.__enter__.return_value = cursor
        self.dbc_instance = DBSAccess(db_con)
        self.assertEqual(self.dbc_instance.get_db_perm("dbc"), 500)

    def test_get_db_free_perm_expects_None(self):
        db_con = MagicMock(UdaExecConnection)
        cursor = MagicMock(UdaExecCursor)
        cursor.fetchone.return_value = None
        db_con.cursor.return_value.__enter__.return_value = cursor
        self.dbc_instance = DBSAccess(db_con)
        self.assertEqual(self.dbc_instance.get_db_perm("dbc"), "dbc does not exist")