Unit testing in python - Can we mock an SQL query result in python?

user_12100101 0
Viewed 4095

Summary : I have a method which i would like to test. This methods involves connecting the database,extracting the result (from sql query) and saving it into a data frame.

I would like to know, is there a way I can mock the result retrieved from the SQL query?

I tried unittest library, created a dataframe like the one received from the SQl query. But I am unable to figure out how can I pass this variable to the method while testing.

Please share your insights. My method is something like :

    try:
        db_connect = sqlite3.connect("address_zip.db")
        db_cur = db_connect.cursor()
        db_cur.execute(""" Select STD_ADDRESS,STD_CITY,STD_STATE from address""")
        sql_result=pd.read_sql_query(""" Select STD_ADDRESS,STD_CITY,STD_STATE from address""",db_connect)
        df = pd.DataFrame(sql_result, columns=['STD_ADDRESS', 'STD_CITY', 'STD_STATE'])
        print(f"The data from database is :\n {df}")
        db_cur.close()
        db_connect.close()
    except Exception as e:
        print(f"DB exception: {e}")

        return df

There are several actions done on the variable "df" in the original code, my agenda is to mock this object df.

But I am unable to understand how? Sorry this may be very basic question, I am new to python.

python-unittest.mock unit-testing python
1 Answers

Let's suggest the code to test would be

code_to_test.py

import sqlite3
import pandas as pd


def function_to_tets():
    df = None
    try:
        db_connect = sqlite3.connect("address_zip.db")
        db_cur = db_connect.cursor()
        db_cur.execute(""" Select STD_ADDRESS,STD_CITY,STD_STATE from address""")
        sql_result = pd.read_sql_query(""" Select STD_ADDRESS,STD_CITY,STD_STATE from address""", db_connect)
        df = pd.DataFrame(sql_result, columns=['STD_ADDRESS', 'STD_CITY', 'STD_STATE'])
        print(f"The data from database is :\n {df}")
        db_cur.close()
        db_connect.close()
    except Exception as e:
        print(f"DB exception: {e}")

    return df

So, the simple solution is to mock any call chain from sqlite3 and mock pd.read_sql_query's return value

tests/test_sql.py

from code_to_test import function_to_tets
from unittest import mock
import pandas as pd


@mock.patch('code_to_test.sqlite3')
@mock.patch('code_to_test.pd.read_sql_query')
def test_sql_query(read_sql_query_mock, sqlite_mock):
    read_sql_query_mock.return_value = pd.DataFrame({'STD_ADDRESS': ['address 1'],
                                                     'STD_CITY': ['city 1'],
                                                     'STD_STATE': ['state 1']})
    assert function_to_tets().to_dict(orient='list') == {'STD_ADDRESS': ['address 1'],
                                                         'STD_CITY': ['city 1'],
                                                         'STD_STATE': ['state 1']}
user_74170700
Related