TST add more test for discovery module (#80)

scikit-learn-contrib · May 1, 2024 · d583456 · d583456
1 parent c0e26ed
commit d583456
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/skltemplate/utils/tests/test_discovery.py b/skltemplate/utils/tests/test_discovery.py
@@ -1,13 +1,25 @@
 # Authors: scikit-learn-contrib developers
 # License: BSD 3 clause
 
+import pytest
+
 from skltemplate.utils.discovery import all_displays, all_estimators, all_functions
 
 
 def test_all_estimators():
     estimators = all_estimators()
     assert len(estimators) == 3
 
+    estimators = all_estimators(type_filter="classifier")
+    assert len(estimators) == 1
+
+    estimators = all_estimators(type_filter=["classifier", "transformer"])
+    assert len(estimators) == 2
+
+    err_msg = "Parameter type_filter must be"
+    with pytest.raises(ValueError, match=err_msg):
+        all_estimators(type_filter="xxxx")
+
 
 def test_all_displays():
     displays = all_displays()