Tuesday, 2 February 2021

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

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. Could someone help?

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 but not sure.

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)


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

No comments:

Post a Comment