class TestDataset(unittest.TestCase): def setUp(self): self._data = {'col1': [1, 2], 'col2': [3, 4], 'col4': [5, 6]} self._dataframe = pd.DataFrame(data=self._data) self._dataset = DataModel(self._dataframe) def test_validate_columns_invalid(self): with self.assertRaises(RuntimeError): self._dataset.validate_columns(['col3']) def test_validate_columns(self): self._dataset.validate_columns(['col1']) def test_feature_columns(self): intended_columns = ['col1', 'col2'] self._dataset.set_feature_columns(intended_columns) feature_columns = self._dataset.get_feature_columns() result_columns = list(feature_columns.columns.values) self.assertEqual(result_columns, intended_columns) def test_target_column(self): intended_column = 'col1' self._dataset.set_target_column(intended_column) target_column = self._dataset.get_target_column() self.assertEqual(target_column.tolist(), self._data[intended_column])
def setUp(self): data_array = {'feat_A': [1, 2, 3], 'feat_B': [8, 6, 4], 'target': [9, 8, 7]} df2 = df = pd.DataFrame(data_array) train_model = DataModel(df) train_model.set_target_column('target') train_model.set_feature_columns(['feat_A', 'feat_B']) eval_model = DataModel(df2) eval_model.set_target_column('target') eval_model.set_feature_columns(['feat_A', 'feat_B']) self.arti = AI('test', 'test/test') self.arti.training_data = train_model self.arti.evaluation_data = eval_model