python - 使用 pd.read_sql 模拟多个数据库调用

Question

我有一个函数使用 pd.read_sql直接从 Python 进行两个 SQL 查询，如下所示:

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

当我从数据库调用时，这很好用。
我现在想模拟这些数据库调用以进行测试，以便当 source_data_query运行，而不是调用数据库，它读入一个测试数据帧。 other_source_data_query 类似.
更新:按照下面亚伦的建议进行编辑:

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))

我收到以下错误: FAILED tests/unit_tests/functions/test.py::GetSQLQueriesTest::test_get_sql_queries - ModuleNotFoundError: No module named 'pd'我肯定在我的环境中安装了 Pandas 。

Answer 1

使用 side_effect 模拟 pd.read_sql 的返回值基于 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))