def test_add_data_single_column(self):
"""Should add a single column of data to the Histogram"""
hist = Histogram()
test_df = self.create_test_df()
column_to_ad = test_df.select(F.col('value'))
hist.add_data(column_to_ad)
self.assertEqual(1, len(hist.col_list))
def test_add_data_list_of_columns(self):
"""Should add all columns from the list of columns to the Histogram"""
test_df = self.create_test_df()
column_to_ad = test_df.select(F.col('value'))
column_to_ad_2 = test_df.select(F.col('value2'))
hist = Histogram()
hist.add_data([column_to_ad, column_to_ad_2])
self.assertEqual(2, len(hist.col_list))
def test_add_data_entire_dataframe(self):
"""Should add all columns of a dataframe to the histogram"""
test_df = self.create_test_df()
hist = Histogram()
hist.add_data(test_df)
self.assertEqual(2, len(hist.col_list))