将“assertTrue”单元测试应用于pandas DataFrame？

Question

使用python的单元测试，如何检查pandas DataFrame df的列名是否与列表的列名一致？

import unittest

column_names_list = ["col1","col2","col3"]

assertTrue(df.columns in column_names_list)

例如，我应该为此df：

获取False

col1   col2   col3   col4
aasa   4546   5643   shsh
...

Answer 1

使用set，您可以找到两个列表之间的差异，assert表示它是0

In [40]: assert len(set(df.columns) - set(['A', 'B', 'C'])) == 0

In [41]: assert len(set(df.columns) - set(['A', 'B', 'D'])) == 0
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
----> 1 assert len(set(df.columns) - set(['A', 'B', 'D'])) == 0

Answer 2

您也可以使用集合的相等性。这样一个断言错误可能更有用（使用大括号的大括号）：

匹配

df = pd.DataFrame([], columns=['col1', 'col2', 'col3'])
assert set(df.columns) == {'col1', 'col2', 'col3'}

不匹配

df = pd.DataFrame([], columns=['col1', 'col2'])
assert set(df.columns) == {'col1', 'col2', 'col3'}