How to use a pandas data frame in a unit test

Question

I am developing a set of python scripts to pre-process a dataset then produce a series of machine learning models using scikit-learn. I would like to develop a set of unittests to check the data pre-processing functions, and would like to be able to use a small test pandas dataframe for which I can determine the answers for and use it in assert statements.

I cannot seem to get it to load the dataframe and to pass it to the unit tests using self. My code looks something like this;

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() 

Thanks for the help.

Answer 1

Pandas has some utilities for testing.

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) 

Answer 2

If you are using latest pandas, I think the following way is a bit cleaner:

import pandas as pd  pd.testing.assert_frame_equal(my_df, expected_df) pd.testing.assert_series_equal(my_series, expected_series) pd.testing.assert_index_equal(my_index, expected_index) 

Each of these functions will raise AssertionError if they are not "equal".

For more information and options: https://pandas.pydata.org/pandas-docs/stable/reference/general_utility_functions.html#testing-functions