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