updated model names

Author: anwurl

tests/workflows/test_octo_t2e.py

@@ -6,6 +6,7 @@
 import pandas as pd
 import pytest
 
+from octopus.models import ModelName
 from octopus.modules import Octo
 from octopus.study import OctoTimeToEvent
 
@@ -93,7 +94,7 @@ def test_octo_task_configuration(self):
             task_id=0,
             depends_on=None,
             description="step_1",
-            models=["CatBoostCoxSurvival"],
+            models=[ModelName.CatBoostCoxSurvival],
             n_trials=12,
             max_features=6,
             ensemble_selection=True,
@@ -108,9 +109,9 @@ def test_octo_task_configuration(self):
         assert octo_task.max_features == 6
         assert octo_task.ensemble_selection is True
         assert octo_task.ensel_n_save_trials == 10
-        assert octo_task.models == ["CatBoostCoxSurvival"]
+        assert octo_task.models == [ModelName.CatBoostCoxSurvival]
 
-    @pytest.mark.parametrize("model_name", ["CatBoostCoxSurvival", "XGBoostCoxSurvival"])
+    @pytest.mark.parametrize("model_name", [ModelName.CatBoostCoxSurvival, ModelName.XGBoostCoxSurvival])
     def test_single_model_configuration(self, model_name):
         """Test configuration with each survival model individually."""
         octo_task = Octo(
@@ -136,22 +137,22 @@ def test_multi_model_configuration(self):
             task_id=0,
             depends_on=None,
             description="step_1",
-            models=["CatBoostCoxSurvival", "XGBoostCoxSurvival"],
+            models=[ModelName.CatBoostCoxSurvival, ModelName.XGBoostCoxSurvival],
             n_trials=12,
             max_features=6,
             ensemble_selection=True,
             ensel_n_save_trials=10,
         )
 
-        assert {str(m) for m in octo_task.models} == {"CatBoostCoxSurvival", "XGBoostCoxSurvival"}
+        assert {str(m) for m in octo_task.models} == {ModelName.CatBoostCoxSurvival, ModelName.XGBoostCoxSurvival}
 
     def test_ensemble_selection_configuration(self):
         """Test ensemble selection configuration."""
         octo_task = Octo(
             task_id=0,
             depends_on=None,
             description="step_1",
-            models=["CatBoostCoxSurvival"],
+            models=[ModelName.CatBoostCoxSurvival],
             n_trials=12,
             max_features=6,
             ensemble_selection=True,
@@ -167,7 +168,7 @@ def test_hyperparameter_optimization_configuration(self):
             task_id=0,
             depends_on=None,
             description="step_1",
-            models=["CatBoostCoxSurvival"],
+            models=[ModelName.CatBoostCoxSurvival],
             n_trials=12,
             max_features=6,
             ensemble_selection=True,
@@ -177,7 +178,7 @@ def test_hyperparameter_optimization_configuration(self):
             penalty_factor=1.5,
         )
 
-        assert "CatBoostCoxSurvival" in octo_task.models
+        assert ModelName.CatBoostCoxSurvival in octo_task.models
         assert octo_task.n_trials == 12
         assert octo_task.max_features == 6
         assert octo_task.ensemble_selection is True
@@ -211,7 +212,7 @@ def test_octo_timetoevent_actual_execution(self, survival_dataset):
                         task_id=0,
                         depends_on=None,
                         description="step_1",
-                        models=["CatBoostCoxSurvival"],
+                        models=[ModelName.CatBoostCoxSurvival],
                         n_trials=12,
                         max_features=6,
                         ensemble_selection=True,
@@ -258,7 +259,7 @@ def test_full_configuration_parameters(self):
             task_id=0,
             depends_on=None,
             description="step_1",
-            models=["CatBoostCoxSurvival"],
+            models=[ModelName.CatBoostCoxSurvival],
             n_trials=12,
             max_features=6,
             ensemble_selection=True,
@@ -279,7 +280,7 @@ def test_full_configuration_parameters(self):
         assert octo_task.task_id == 0
         assert octo_task.depends_on is None
         assert octo_task.description == "step_1"
-        assert octo_task.models == ["CatBoostCoxSurvival"]
+        assert octo_task.models == [ModelName.CatBoostCoxSurvival]
         assert octo_task.model_seed == 0
         assert octo_task.n_jobs == 1
         assert octo_task.max_outl == 0