updated blend_models test and create test_stack_models.py

ajinkyapuar · May 9, 2020 · b7f81d6 · b7f81d6
1 parent 2525193
commit b7f81d6
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/pycaret/tests/test_blend_models.py b/pycaret/tests/test_blend_models.py
@@ -1,12 +1,12 @@
 import os, sys
 sys.path.insert(0, os.path.abspath(".."))
 
-#compare_models_test
+#blend_models_test
 import pytest
 import pycaret.regression
 import pycaret.datasets
 
-def test_compare_models():
+def test_blend_models():
     data = pycaret.datasets.get_data('boston')
     data = data.head(50)
     reg1 = pycaret.regression.setup(data, target='medv',silent=True, html=False, session_id=123)

diff --git a/pycaret/tests/test_stack_models.py b/pycaret/tests/test_stack_models.py
@@ -0,0 +1,15 @@
+import os, sys
+sys.path.insert(0, os.path.abspath(".."))
+
+#stack_models_test
+import pytest
+import pycaret.regression
+import pycaret.datasets
+
+def test_stack_models():
+    data = pycaret.datasets.get_data('boston')
+    data = data.head(50)
+    reg1 = pycaret.regression.setup(data, target='medv',silent=True, html=False, session_id=123)
+    estimator_list = pycaret.regression.compare_models(blacklist = ['catboost', 'tr'], n_select=3, verbose=False) #select top 3
+    stacker = pycaret.regression.stack_models(estimator_list=estimator_list, improve_only=False, verbose=False)
+    assert type(stacker) == list