fix tests

EpistasisLab · Jan 12, 2021 · 9a2e2fb · 9a2e2fb
1 parent 55dc719
commit 9a2e2fb
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/tests/tpot_tests.py b/tests/tpot_tests.py
@@ -833,8 +833,8 @@ def test_TPOT_ensemble_1():
         ensemble=True
     )
     tpot_obj.fit(training_features, training_target)
-    print(tpot_obj.fitted_pipeline_)
-
+    from sklearn.ensemble import VotingClassifier
+    assert isinstance(tpot_obj.fitted_pipeline_, VotingClassifier)
 
 def test_TPOT_ensemble_2():
     """Assert that TPOT uses ensemble methods."""
@@ -849,7 +849,8 @@ def test_TPOT_ensemble_2():
         ensemble=True
     )
     tpot_obj.fit(training_features, training_target)
-    print(tpot_obj.fitted_pipeline_)
+    from sklearn.ensemble import VotingRegressor
+    assert isinstance(tpot_obj.fitted_pipeline_, VotingRegressor)
 
 def test_fit_GroupKFold():
     """Assert that TPOT properly handles the group parameter when using GroupKFold."""