Python Unit Test : How do I test loading data from database?-CodePudding

I'm using psycopg2-binary to connect to a database. This is the class I've written to query data from the database given a query and return as a dataframe. I'm stuck on how I would unit test this function without having to hit the real database.

code :

import pandas as pd

class HitDB(object):
    def __init__(self, db_con):
        self.conn = db_con
        

    def load_data(self, query):
        df = pd.read_sql(
            sql = query, 
            con = self.conn
        )
        return df

CodePudding user response：

Since all the function does is pass its query and self.conn to another function, mock pd.read_sql and check it receives the expected arguments.

Which suggests getting rid of this function and just call pd.read_sql directly.

If it did something more significant with the database, run a local PostgreSQL server (locally or with something like Docker), load the schema, and use factories to populate it with test data as needed.