如何为读取Parquet文件返回DataFrame的Python函数编写单元测试？

针对你的read_data函数，有两种实用的单元测试思路，不用把流程搞复杂：

一、用真实测试Parquet文件（推荐，贴近实际场景） #

这种方式直接模拟真实的文件读取场景，能验证Parquet文件的读写逻辑是否正常，步骤很简单：

1. 生成临时测试文件：用pandas创建包含目标列的小DataFrame，写入临时目录的Parquet文件；
2. 调用函数并验证结果：传入临时目录路径，检查返回的DataFrame是否和预期一致；
3. 自动清理临时文件：用Python的tempfile模块自动管理临时目录，测试完成后自动删除。

示例代码（用pytest） #

import pytest
import pandas as pd
import pyarrow as pq
from tempfile import TemporaryDirectory
from your_module import read_data  # 替换成你的函数所在模块名

def test_read_data_with_real_parquet():
    # 准备测试数据
    test_data = pd.DataFrame({
        'userID': [1, 2, 3],
        'itemID': [101, 102, 103],
        'action': ['click', 'view', 'purchase'],
        'extra_col': ['ignore', 'this', 'column']  # 模拟多余列
    })
    
    # 写入临时Parquet文件
    with TemporaryDirectory() as temp_dir:
        pq.write_table(pq.Table.from_pandas(test_data), f"{temp_dir}/test_data.parquet")
        
        # 调用函数并验证
        result_df = read_data(temp_dir)
        
        # 验证列是否正确
        assert list(result_df.columns) == ['userID', 'itemID', 'action']
        # 验证数据内容匹配
        pd.testing.assert_frame_equal(result_df, test_data[['userID', 'itemID', 'action']])

二、Mock Parquet读取（适合快速测试逻辑） #

如果不想生成物理文件，或者只需要测试函数的列选择逻辑，可以Mockpyarrow.to_pandas方法，让它直接返回预设的DataFrame，跳过真实的文件IO。

示例代码（用pytest-mock） #

import pytest
import pandas as pd
from unittest.mock import patch
from your_module import read_data

def test_read_data_with_mock():
    # 准备模拟返回的DataFrame
    mock_data = pd.DataFrame({
        'userID': [1, 2],
        'itemID': [101, 102],
        'action': ['click', 'view'],
        'other_col': [456, 789]
    })
    
    # Mock pyarrow.to_pandas方法
    with patch('pyarrow.to_pandas') as mock_to_pandas:
        mock_to_pandas.return_value = mock_data
        
        # 调用函数并验证
        result_df = read_data('/fake/dir')
        
        # 验证mock被正确调用
        mock_to_pandas.assert_called_once_with('/fake/dir')
        # 验证列选择正确
        assert list(result_df.columns) == ['userID', 'itemID', 'action']
        pd.testing.assert_frame_equal(result_df, mock_data[['userID', 'itemID', 'action']])

额外测试场景建议 #

• 测试目录无文件：验证函数是否抛出合理的异常（比如FileNotFoundError）；
• 测试文件缺少指定列：模拟只有userID和itemID的Parquet文件，验证函数是否抛出KeyError；
• 测试多Parquet文件：在临时目录下写入多个Parquet文件，验证函数是否能正确合并所有数据。

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