将“ assertTrue”单元测试应用于熊猫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'}