如何对生成CSV文件的函数进行单元测试（含DataFrame近似断言）

在编写这类单元测试时，核心是利用pandas自带的断言工具或结合numpy的数值比较逻辑，处理浮点数的精度误差问题，以下是几种实用方案：

1. 使用pandas的assert_frame_equal（推荐） #

pandas专门提供了pandas.testing.assert_frame_equal函数，支持通过容差参数处理近似相等的情况，这是最直接的方案：

代码示例 #

import pandas as pd
from pandas.testing import assert_frame_equal

def test_generate_csv():
    # 1. 调用你的生成函数，读取输出CSV为DataFrame
    generate_your_csv_function(output_path="output.csv")
    df_actual = pd.read_csv("output.csv")
    
    # 2. 读取预期的基准CSV为DataFrame
    df_expected = pd.read_csv("expected.csv")
    
    # 3. 执行近似相等断言
    assert_frame_equal(
        df_actual, 
        df_expected,
        atol=1e-5,  # 绝对容差：允许的最大绝对差异，比如8.0和7.99999会被判定相等
        rtol=1e-3,  # 相对容差：允许的最大相对差异，适合不同量级的数值
        check_dtype=False,  # 可选：忽略int和float这类细微类型差异
        check_like=True  # 可选：忽略索引/列的顺序差异（如果不需要严格顺序）
    )

参数说明 #

• atol：绝对容差，比如设置为1e-5时，8.0和7.99999会被判定为相等
• rtol：相对容差，计算公式为|a - b| <= (atol + rtol * |b|)，适合处理不同量级的数值
• 若只需检查部分列，可通过subset参数指定列名列表，比如subset=["col1", "col2"]

2. 结合numpy的np.allclose自定义断言 #

如果需要更灵活的控制（比如针对不同列设置不同容差），可以用numpy的allclose函数单独验证数值列：

代码示例 #

import pandas as pd
import numpy as np

def test_generate_csv():
    df_actual = pd.read_csv("output.csv")
    df_expected = pd.read_csv("expected.csv")
    
    # 先确保列和索引结构一致
    assert df_actual.columns.tolist() == df_expected.columns.tolist()
    assert df_actual.index.tolist() == df_expected.index.tolist()
    
    # 遍历数值列，逐一验证近似相等
    for col in df_actual.select_dtypes(include=np.number).columns:
        # 针对特定列可单独调整容差
        tol = 1e-4 if col == "precision_col" else 1e-5
        assert np.allclose(df_actual[col], df_expected[col], atol=tol), f"列{col}数值不匹配"
    
    # 非数值列严格匹配
    for col in df_actual.select_dtypes(exclude=np.number).columns:
        assert df_actual[col].equals(df_expected[col]), f"列{col}内容不匹配"

3. 处理特殊场景 #

• 若CSV读取时出现浮点精度问题，可在pd.read_csv中指定dtype参数统一类型，避免不必要的差异
• 对于时间列或其他非数值列，确保读取时的解析规则一致（比如使用parse_dates参数统一解析格式）

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