Python单元测试：如何测试从数据库加载数据的功能？

你的HitDB类设计已经具备了依赖注入的优势（通过构造函数传入数据库连接），这让单元测试变得简单——只需要用Mock工具模拟依赖，不需要真实数据库。下面是两种实用的测试方案：

方案1：直接Mock pandas.read_sql #

load_data的核心逻辑就是调用pd.read_sql并返回结果，所以直接Mock这个函数，让它返回预设的测试数据即可，完全不需要关心数据库连接的细节。

测试代码（基于unittest框架） #

import unittest
from unittest.mock import patch
import pandas as pd
from your_module import HitDB  # 替换为你的HitDB所在模块名

class TestHitDBLoadData(unittest.TestCase):
    def test_load_data_returns_correct_dataframe(self):
        # 构造测试用的假数据集
        expected_df = pd.DataFrame({
            'id': [101, 102, 103],
            'name': ['Alice', 'Bob', 'Charlie']
        })
        
        # 模拟pd.read_sql，让它返回预设的DataFrame
        with patch('pandas.read_sql') as mock_read_sql:
            mock_read_sql.return_value = expected_df
            
            # 传入任意假连接（不会被真实调用）
            hit_db = HitDB(db_con='fake_db_connection')
            # 调用待测试方法
            result_df = hit_db.load_data('SELECT id, name FROM users')
            
            # 验证返回结果是否符合预期
            pd.testing.assert_frame_equal(result_df, expected_df)
            # 验证pd.read_sql是否被正确调用（参数匹配）
            mock_read_sql.assert_called_once_with(
                sql='SELECT id, name FROM users',
                con='fake_db_connection'
            )

方案2：Mock数据库连接对象（更贴近真实调用链） #

如果你想模拟完整的连接调用流程（虽然对这个方法来说没必要，但可以用于更复杂的场景），可以用MagicMock创建一个模拟的数据库连接对象，确保pd.read_sql能正确接收并使用它。

测试代码 #

import unittest
from unittest.mock import MagicMock, patch
import pandas as pd
from your_module import HitDB

class TestHitDBLoadDataWithMockedConn(unittest.TestCase):
    def test_load_data_uses_db_connection_correctly(self):
        test_df = pd.DataFrame({
            'product_id': [1, 2],
            'price': [99.9, 199.9]
        })
        
        # 创建模拟的数据库连接
        mock_conn = MagicMock()
        
        # 模拟pd.read_sql行为
        with patch('pandas.read_sql', return_value=test_df):
            hit_db = HitDB(db_con=mock_conn)
            result = hit_db.load_data('SELECT product_id, price FROM products')
            
            # 验证结果正确性
            pd.testing.assert_frame_equal(result, test_df)
            # 验证连接对象被正确传入read_sql
            pd.read_sql.assert_called_once_with(
                sql='SELECT product_id, price FROM products',
                con=mock_conn
            )

关键说明 #

• 我们测试的是HitDB的逻辑：即load_data是否正确调用pd.read_sql并返回结果，而不是测试数据库连接、SQL语法正确性（这些属于集成测试范畴）。
• 依赖注入的设计是关键：因为HitDB不自己创建连接，而是通过构造函数接收，所以我们可以轻松传入模拟对象替代真实连接。
• 可以扩展测试场景：比如测试空结果、异常情况（模拟pd.read_sql抛出异常，验证load_data是否正确传递异常）。

内容的提问来源于stack exchange，提问作者ACCCCCCC