如何编写单元测试验证DataFrame包含指定预期列？

下面针对你的需求，提供两种常用Python测试框架的实现方案，直接套用即可：

方案1：使用pytest框架 #

import pandas as pd
import pytest

def test_expected_columns_exist():
    # 初始化测试数据
    d = {'ID_EMPLOYEE': [12, 35, 56, 46], 'Number':[0,1,2,30], 'Location_EMPLOYEE':["US","US","Austria","France"], 'Salary':[100,200,100,160]}
    df_output = pd.DataFrame(d)
    List_Match = ["Location_EMPLOYEE","ID_EMPLOYEE"]
    
    # 方式A：简洁的集合子集验证（一步确认所有列存在）
    assert set(List_Match).issubset(df_output.columns)
    
    # 方式B：逐个列验证（失败时能明确显示缺失的列）
    for col in List_Match:
        assert col in df_output.columns, f"预期列 {col} 未在DataFrame中找到"

运行测试时，执行pytest 你的测试文件名.py -v即可，失败时会给出清晰的错误提示。

方案2：使用unittest框架 #

import pandas as pd
import unittest

class TestDFColumns(unittest.TestCase):
    def setUp(self):
        # 提前初始化测试数据，每个测试方法都会复用
        d = {'ID_EMPLOYEE': [12, 35, 56, 46], 'Number':[0,1,2,30], 'Location_EMPLOYEE':["US","US","Austria","France"], 'Salary':[100,200,100,160]}
        self.df_output = pd.DataFrame(d)
        self.List_Match = ["Location_EMPLOYEE","ID_EMPLOYEE"]
    
    def test_all_columns_present(self):
        # 集合子集验证
        self.assertTrue(set(self.List_Match).issubset(self.df_output.columns))
    
    def test_each_column_exists(self):
        # 逐个验证列，subTest能让所有列的验证都执行完再报告失败
        for col in self.List_Match:
            with self.subTest(col=col):
                self.assertIn(col, self.df_output.columns)

if __name__ == '__main__':
    unittest.main()

运行测试直接执行python 你的测试文件名.py，subTest的作用是即使某一列缺失，也会继续检查其他列，方便一次性定位所有问题。

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