add unit tests for predict speed improvements (#312)

tests/tests_algorithms/test_baseline.py

@@ -82,3 +82,9 @@ def test_save_and_load(self):
         y_predicted = dt2.predict(self.X)
         loss2 = metric(self.y, y_predicted)
         assert_almost_equal(loss, loss2)
+
+    def test_is_fitted(self):
+        model = BaselineRegressorAlgorithm({"ml_task": "regression"})
+        self.assertFalse(model.is_fitted())
+        model.fit(self.X, self.y)
+        self.assertTrue(model.is_fitted())

tests/tests_algorithms/test_catboost.py

@@ -144,3 +144,9 @@ def test_get_metric_name(self):
         params["loss_function"] = "MultiClass"
         model = CatBoostAlgorithm(params)
         self.assertEqual(model.get_metric_name(), "logloss")
+
+    def test_is_fitted(self):
+        cat = CatBoostAlgorithm(self.params)
+        self.assertFalse(cat.is_fitted())
+        cat.fit(self.X, self.y)
+        self.assertTrue(cat.is_fitted())

tests/tests_algorithms/test_decision_tree.py

@@ -60,3 +60,10 @@ def test_save_and_load(self):
 
         # Finished with temp file, delete it
         os.remove(filename)
+
+    def test_is_fitted(self):
+        params = {"max_depth": 1, "seed": 1, "ml_task": "regression"}
+        model = DecisionTreeRegressorAlgorithm(params)
+        self.assertFalse(model.is_fitted())
+        model.fit(self.X, self.y)
+        self.assertTrue(model.is_fitted())

tests/tests_algorithms/test_extra_trees.py

@@ -114,3 +114,10 @@ def test_save_and_load(self):
         y_predicted = rf2.predict(self.X)
         loss2 = metric(self.y, y_predicted)
         assert_almost_equal(loss, loss2)
+
+    def test_is_fitted(self):
+        params = {"trees_in_step": 50, "ml_task": "binary_classification"}
+        model = ExtraTreesAlgorithm(params)
+        self.assertFalse(model.is_fitted())
+        model.fit(self.X, self.y)
+        self.assertTrue(model.is_fitted())

tests/tests_algorithms/test_knn.py

@@ -70,3 +70,11 @@ def test_fit_predict(self):
         la.fit(self.X, self.y)
         y_predicted = la.predict(self.X)
         self.assertTrue(metric(self.y, y_predicted) < 0.6)
+
+
+    def test_is_fitted(self):
+        params = {"ml_task": "binary_classification"}
+        model = KNeighborsAlgorithm(params)
+        self.assertFalse(model.is_fitted())
+        model.fit(self.X, self.y)
+        self.assertTrue(model.is_fitted())

tests/tests_algorithms/test_lightgbm.py

@@ -120,6 +120,8 @@ def test_restricted_characters_in_feature_name(self):
         lgb.fit(X, y)
         lgb.predict(X)
 
-
-if __name__ == "__main__":
-    unittest.main()
+    def test_is_fitted(self):
+        model = LightgbmAlgorithm(self.params)
+        self.assertFalse(model.is_fitted())
+        model.fit(self.X, self.y)
+        self.assertTrue(model.is_fitted())

tests/tests_algorithms/test_linear.py

@@ -104,3 +104,9 @@ def test_save_and_load(self):
         y_predicted = model2.predict(self.X)
         loss2 = metric(self.y, y_predicted)
         assert_almost_equal(loss, loss2)
+
+    def test_is_fitted(self):
+        model = LinearAlgorithm({"ml_task": "binary_classification"})
+        self.assertFalse(model.is_fitted())
+        model.fit(self.X, self.y)
+        self.assertTrue(model.is_fitted())

tests/tests_algorithms/test_nn.py

@@ -152,3 +152,9 @@ def test_fit_predict(self):
         y_predicted = nn.predict(self.X)
         loss = metric(self.y, y_predicted)
         self.assertLess(loss, 2)
+
+    def test_is_fitted(self):
+        model = MLPAlgorithm(self.params)
+        self.assertFalse(model.is_fitted())
+        model.fit(self.X, self.y)
+        self.assertTrue(model.is_fitted())

tests/tests_algorithms/test_random_forest.py

@@ -114,3 +114,9 @@ def test_save_and_load(self):
         y_predicted = rf2.predict(self.X)
         loss2 = metric(self.y, y_predicted)
         assert_almost_equal(loss, loss2)
+
+    def test_is_fitted(self):
+        model = RandomForestAlgorithm({"ml_task": "binary_classification"})
+        self.assertFalse(model.is_fitted())
+        model.fit(self.X, self.y)
+        self.assertTrue(model.is_fitted())

tests/tests_algorithms/test_xgboost.py

@@ -109,3 +109,10 @@ def test_get_metric_name(self):
         params = {"eval_metric": "rmse"}
         model = XgbAlgorithm(params)
         self.assertEqual(model.get_metric_name(), "rmse")
+
+    def test_is_fitted(self):
+        params = {"objective": "binary:logistic", "eval_metric": "logloss"}
+        model = XgbAlgorithm(params)
+        self.assertFalse(model.is_fitted())
+        model.fit(self.X, self.y)
+        self.assertTrue(model.is_fitted())