def test_prediction_default(self): """Test predict method of decision tree classifier""" classifier = DecisionTreeClassifier() predict_field = "two" classifier.fit(self.data1_df, predict_field) predict_input = pd.Series(data=[5., ], index=["one"]) prediction = classifier.predict(predict_input) self.assertEqual(prediction, self.data1_df[predict_field].mode()[0])
def test_prediction_lookup(self): """Test predict method of decision tree classifier""" classifier = DecisionTreeClassifier() predict_field = "two" classifier.fit(self.data1_df, predict_field) predict_input = pd.Series(data=[2., ], index=["one", ]) prediction = classifier.predict(predict_input) self.assertEqual(prediction, self.data1_df[self.data1_df["one"] == 2.][predict_field].iloc[0])