使用 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.

同样

更新：根据以下 Aaron 的建议进行编辑：

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

如果您的目标不是阅读整个 table 而只是一个测试样本，您可能希望将 LIMIT [#ofrows] 添加到查询的末尾，但是 [#ofrows] 是您要测试的行数运行.

您还可以使用 ORDER BY RAND() LIMIT [#ofrows]

以随机顺序获取行

如果这不是您的问题，您需要就您要查找的内容提供更多说明（请参阅 post 下的评论）。我还建议清理您问题中的代码，因为看起来您正在关闭 "" 紧接在 other_source_data_query =

之后

Answer 2

根据sql.

使用side_effect模拟pd.read_sql的return值

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

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

Mocking out multiple database calls using pd.read_sql

python

sql

mocking