Pytest中Pandas DataFrame快照测试失败后的技术问询

背景 #

用Pytest测试数据清理流程时，直接用df1.equals(df2)断言失败后，控制台输出的信息很难快速定位不匹配的数据，尤其当DataFrame规模较大时。针对你提出的三个问题，给出具体解决方案：

1. 测试失败时定位不匹配行或生成差异文件 #

直接用equals只能得到布尔结果，我们可以结合Pandas的compare方法或自定义逻辑，精准输出差异：

方案1：打印并保存差异详情

修改测试代码，在断言前先检查差异：

import pandas as pd

def test_data_cleaning_process():
    df1 = pd.read_csv('df1.csv')
    df2 = pd.read_csv('df2.csv')
    df1 = cleanup(df1)

    if not df1.equals(df2):
        # 生成带差异的对比表，只保留不相等的内容
        diff = df1.compare(df2, keep_shape=False, keep_equal=False)
        print("=== DataFrame 差异详情 ===")
        print(diff)
        # 保存差异到CSV文件，方便后续审核
        diff.to_csv('df_diff.csv', index=True)
        # 提取所有不匹配的行索引
        mismatched_rows = diff.dropna(how='all').index.unique()
        print(f"\n不匹配的行索引：{list(mismatched_rows)}")
    
    assert df1.equals(df2)

compare方法会把两个DataFrame的差异按列拆分（显示self和other的值），keep_shape=False会自动过滤完全匹配的行/列，让输出更简洁。

方案2：自定义布尔掩码定位差异

如果你的Pandas版本较低（<1.1.0），可以用掩码手动筛选：

# 找出所有不匹配的单元格
mask = df1 != df2
# 筛选出至少有一个不匹配单元格的行
mismatched_df = df1[mask.any(axis=1)]
print("不匹配的行内容：")
print(mismatched_df)

2. 测试失败时自动覆盖快照文件 #

可以通过两种方式实现：

方案1：手动编写逻辑

在断言失败的异常处理中，将当前处理后的DataFrame覆盖基准文件（注意：建议先确认差异是预期的再开启）：

def test_data_cleaning_process():
    df1 = pd.read_csv('df1.csv')
    df2 = pd.read_csv('df2.csv')
    df1 = cleanup(df1)

    try:
        assert df1.equals(df2)
    except AssertionError:
        # 确认要更新快照时，取消注释下面一行
        # df1.to_csv('df2.csv', index=False)
        print("快照已更新，请验证后提交")
        raise

方案2：使用Pytest插件

用pytest-snapshot插件，它支持自动管理快照文件：

1. 安装插件：pip install pytest-snapshot
2. 修改测试代码使用快照断言：

from pytest_snapshot.plugin import assert_snapshot_match

def test_data_cleaning_process():
    df1 = pd.read_csv('df1.csv')
    df1 = cleanup(df1)
    # 将DataFrame转为字符串格式用于快照对比
    assert_snapshot_match(df1.to_csv(index=False), 'df2_snapshot.csv')

3. 首次运行生成快照，后续测试失败时，用pytest --snapshot-update命令更新快照。

3. 测试数据清理流程的最佳实践 #

• 快照测试优先：维护经过人工验证的基准数据（如df2.csv），作为测试的唯一可信来源，每次重构后对比结果。
• 拆分测试粒度：不要只测最终结果，把清理的每个步骤单独测试（比如处理缺失值、格式转换、过滤无效行等），方便快速定位问题。
• 覆盖边界场景：测试空DataFrame、含异常值（如NaN、特殊字符）的输入、极端规模的数据，确保清理逻辑的鲁棒性。
• 版本控制基准数据：将基准文件（如df2.csv）加入Git等版本控制系统，追踪每次变更的原因，避免误修改。
• 自动生成差异报告：测试失败时自动保存差异文件，无需手动对比整个DataFrame。
• 参数化测试：用@pytest.mark.parametrize测试多组不同的输入数据，覆盖更多业务场景。

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