Merge branch 'main' of https://github.com/xaviernogueira/Tabular_ML

xaviernogueira · Jun 5, 2023 · 405275e · 405275e
2 parents acce427 + 6539476
commit 405275e
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/tests/test_1_factory.py b/tests/test_1_factory.py
@@ -35,3 +35,19 @@ def test_classification_models() -> None:
         classification_models[0],
     )
     assert issubclass(first_model_obj, tabular_ml.base.MLModel)
+
+
+def test_factory_methods() -> None:
+    """Test unregister_model and register_model methods."""
+
+    # get a copy of a model class
+    model_name = tabular_ml.ModelFactory.get_regression_models()[0]
+    model = tabular_ml.ModelFactory.get_regression_model(model_name)
+
+    # unregister the model
+    tabular_ml.ModelFactory.unregister_model(model_name)
+    assert model_name not in tabular_ml.ModelFactory.get_regression_models()
+
+    # re-register the model
+    tabular_ml.ModelFactory.register_model(model)
+    assert model_name in tabular_ml.ModelFactory.get_regression_models()