如何在单元测试中使用PANAS数据框

问题描述

我正在开发一组Python脚本，用于对数据集进行预处理，然后使用SCRKIT-LEARN生成一系列机器学习模型。我想开发一套单元测试来检查数据的预处理功能，并希望能够使用一个小型的 pandas 数据帧来确定答案，并在Assert语句中使用它。

我似乎无法让它加载数据帧并使用Self将其传递给单元测试。我的代码如下所示；

def setUp(self):
    TEST_INPUT_DIR = 'data/'
    test_file_name =  'testdata.csv'
    try:
        data = pd.read_csv(INPUT_DIR + test_file_name,
            sep = ',',
            header = 0)
    except IOError:
        print 'cannot open file'
    self.fixture = data

def tearDown(self):
    del self.fixture

def test1(self):    
    self.assertEqual(somefunction(self.fixture), somevalue)

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

谢谢您的帮助。

---

解决方案

pandas 有一些用于测试的实用程序。

import unittest
import pandas as pd
from pandas.util.testing import assert_frame_equal # <-- for testing dataframes

class DFTests(unittest.TestCase):

    """ class for running unittests """

    def setUp(self):
        """ Your setUp """
        TEST_INPUT_DIR = 'data/'
        test_file_name =  'testdata.csv'
        try:
            data = pd.read_csv(INPUT_DIR + test_file_name,
                sep = ',',
                header = 0)
        except IOError:
            print 'cannot open file'
        self.fixture = data

    def test_dataFrame_constructedAsExpected(self):
        """ Test that the dataframe read in equals what you expect"""
        foo = pd.DataFrame()
        assert_frame_equal(self.fixture, foo)