如何对 dataframe 操作进行单元测试？

Question

我有一个 function：

def example(df: DataFrame) -> DataFrame:
    
    return (df - df.mean()) / df.std()

例如一个测试用例可以是：我创建一个 dataframe：

df = pd.DataFrame(np.array([[1], [4], [7]]),
                   columns=['a'])
print(df)
   a
0  1
1  4
2  7

print(example(df))
     a
0 -1.0
1  0.0
2  1.0

这是我尝试过的：

class testExample(unittest.TestCase):
    def test_valid_inputs(self):
        test_df_1 = pd.DataFrame({'test': [1, 4, 7]})
        expected_result_1 = pd.Series([-1.0, 0.0, 1.0])
        pd.testing.assert_series_equal(xxx.example(test_df_1), expected_result_1, check_names=False)

这给了我错误AssertionError: Series Expected type <class 'pandas.core.series.Series'>, found <class 'pandas.core.frame.DataFrame'> instead ，我尝试了其他一些不同的东西，但没有工作......

Answer 1

主要是口味问题，但我更喜欢使用特殊的setUp方法来初始化数据。 然后可以通过一种以上的测试方法共享它。 在这里我会使用：

class TestExample(unittest.TestCase):
    def setUp(self):
        """Initialize a dataframe for the tests"""
        self.df = pd.DataFrame(np.array([[1], [4], [7]]),
                               columns=['a'])

    def testExample(self):
        """Control the result of the example function"""
        result = example(self.df)
        self.assertEqual([-1.0, 0, 1.0], result['a'].to_list())

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