How to unit test a method that contains a database call in Python

Question

I want to unit test a method that includes a database (SQL Server) call.

I don't want the test to connect to the actual database.

I use unittest for testing, I have done some research and it seems that Mocking could do the trick but not sure about the syntax.

The select statement on the code below returns some integers. I guess that mocking will target the "cursor.execute" and "cursor.fetchall()" parts of the code.

from databaselibrary.Db import Db

class RandomClass():

    def __init__(self, database):
        self.database = database # Main DB for inserting data
            
    def check_file_status(self, trimmed_file_data, file_date):
        cursor = self.database.cursor()

        cursor.execute(f"""SELECT DISTINCT query_id
                FROM wordcloud_count
                WHERE date = '{file_date}'""")
        
        queries_in_DB = set(row.query_id for row in cursor.fetchall())

        queries_in_file = set(trimmed_file_data.keys())

        if queries_in_DB == queries_in_file:
            return False
        
        return True
    
    def run(self):
        print("Hello")


if __name__ == "__main__":
    
    connection_string = 'sql://user:password@server/database'
    
    database = Db(connection_string, autocommit=True)
    
    random = RandomClass(database)
    
    random.run()

The test class could look like that:

import unittest
from unittest.mock import Mock, patch
from project.RandomClass import RandomClass
from datetime import datetime

class testRandomClass(unittest.TestCase):
    
    def setUp(self):
        self.test_class = RandomClass("don't want to put actual database here")
    
    @patch("project.RandomClass.check_file_status",return_value={123, 1234})
    def test_check_file_status(self):
        
        keys = {'1234':'2','123':'1','111':'5'}
        
        result = self.test_class.check_file_status(keys, datetime(1900, 1, 1, 23, 59, 59))
        
        self.assertTrue(result)

user_840582 · Answer

You should mock the db connection object, and the cursor. Then, set the return value of the cursor to return the expected value. I've tested the below code, and used class Row to mock the rows returned from fetchall call:

import unittest
from unittest.mock import MagicMock
from datetime import datetime
from project.RandomClass import RandomClass


class Row(object):
    def __init__(self, x):
        self.query_id = x


class testRandomClass(unittest.TestCase):

    def setUp(self):
        dbc = MagicMock(name="dbconn")
        cursor = MagicMock(name="cursor")
        cursor.fetchall.return_value = [Row(1), Row(2)]
        dbc.cursor.return_value = cursor
        self.test_class = RandomClass(dbc)

    def test_check_file_status(self):
        keys = {'1234': '2', '123': '1', '111': '5'}

        result = self.test_class.check_file_status(keys, datetime(1900, 1, 1, 23, 59, 59))

        self.assertTrue(result)

Since in your RandomClass you iterate rows and get their query_id, you need to use a class (or a named tuple) as the row objects returned by the mock.

You should create the row objects you expected, and set them as the return value of fetchall.

How to unit test a method that contains a database call in Python

1 Answers