Mocking out multiple database calls using pd.read_sql

Question

I have a function which uses pd.read_sql to make two SQL queries directly from Python, as per the below:

def get_sql_queries():
   source_data_query = """
        SELECT
              cf."fund" as 'Fund'
            , cf."Symbol"
   FROM
            sql_table
   """

   transactions = pd.read_sql(
        sql=source_data_query,
        con=DEFAULT_CONNECTION,
    )

   other_source_data_query = """
        SELECT
              cf."fund" as 'Fund'
            , cf."Symbol"
   FROM
            other_sql_table
   """
   marks = pd.read_sql(
        sql=other_source_data_query,
        con=DEFAULT_CONNECTION,
    )
   returns transactions,marks

This works fine when I'm calling from the database.

I would now like to mock out these database calls for testing purposes, such that when source_data_query is run, instead of calling the database, it reads in a test dataframe. Similarly for other_source_data_query.

Update: editing as per Aaron's suggestion below:

import unittest
from unittest import mock
import pandas as pd
from functions import get_transaction_data


class GetSQLQueriesTest(unittest.TestCase):
    @mock.patch('pd.read_sql')
    def test_get_sql_queries(self, mock_read_sql):
        transaction_data = pd.DataFrame(columns=['Fund', 'Symbol'], data=[['Fund 1', 'Symbol 1']])
        mark_data = pd.DataFrame(columns=['Fund', 'Symbol'], data=[['Fund 2', 'Symbol 2']])

        mock_read_sql.side_effect = (transaction_data, mark_data)  # If the order is fixed

        output = get_transaction_data.get_transactions_between_two_dates()
        self.assertEqual(output, (transaction_data, mark_data))

I'm getting the following error:

FAILED tests/unit_tests/functions/test.py::GetSQLQueriesTest::test_get_sql_queries - ModuleNotFoundError: No module named 'pd'

I definitely have pandas installed on my environment.

Answer 1

Use side_effect to mock the return value of pd.read_sql based on sql.

import unittest
from unittest import mock

import pandas as pd

from mymodule import get_sql_queries


class GetSQLQueriesTest(unittest.TestCase):

    @mock.patch('pandas.read_sql')  # Or @mock.patch('mymodule.pd.read_sql')
    def test_get_sql_queries(self, mock_read_sql):
        transaction_data = pd.DataFrame(columns=['Fund', 'Symbol'], data=[['Fund 1', 'Symbol 1']])
        mark_data = pd.DataFrame(columns=['Fund', 'Symbol'], data=[['Fund 2', 'Symbol 2']])

        # mock_read_sql.side_effect = (transaction_data, mark_data)  # If the order is fixed
        mock_read_sql.side_effect = lambda sql, con: (
            transaction_data if ' sql_table' in sql else
            mark_data if 'other_sql_table' in sql else
            None
        )

        output = get_sql_queries()
        self.assertEqual(output, (transaction_data, mark_data))