add test for ExtraTreeRegression

openscilab · Sep 20, 2023 · 4237507 · 4237507
1 parent c90c1f3
commit 4237507
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/tests/test_decision_trees/extra_tree/extra_tree_regression.py b/tests/test_decision_trees/extra_tree/extra_tree_regression.py
@@ -0,0 +1,14 @@
+from sklearn.tree import ExtraTreeRegressor
+
+from pymilo.utils.test_pymilo import pymilo_regression_test
+from pymilo.utils.data_exporter import prepare_simple_regression_datasets
+
+MODEL_NAME = "Extra Tree Regressor"
+
+def extra_tree_regression():
+    x_train, y_train, x_test, y_test = prepare_simple_regression_datasets()
+    # Create Decision Tree Regressor
+    extra_tree_regressor = ExtraTreeRegressor(random_state=0)
+    extra_tree_regressor = extra_tree_regressor.fit(x_train, y_train)
+    assert pymilo_regression_test(
+        extra_tree_regressor, MODEL_NAME, (x_test, y_test)) == True