Fix tests

Author: thomasmeissnercrm

bluecast/ml_modelling/base_classes.py

@@ -284,13 +284,15 @@ def concat_prepare_full_train_datasets(
 
         return x_train, y_train
 
-    def get_early_stopping_callback(self) -> Optional[List[xgb.callback.EarlyStopping]]:
+    def get_early_stopping_callback(
+        self, data_name: str = "test"
+    ) -> Optional[List[xgb.callback.EarlyStopping]]:
         """Create early stopping callback if configured."""
         if self.conf_training.early_stopping_rounds:
             early_stop = xgb.callback.EarlyStopping(
                 rounds=self.conf_training.early_stopping_rounds,
                 metric_name=self.conf_xgboost.xgboost_eval_metric,
-                data_name="test",
+                data_name=data_name,
                 save_best=self.conf_params_xgboost.params["booster"] != "gblinear",
             )
             callbacks = [early_stop]

bluecast/ml_modelling/xgboost.py

@@ -101,8 +101,10 @@ def fit(
         d_train, d_test = self._create_d_matrices(x_train, y_train, x_test, y_test)
         if x_test.empty:
             eval_set = [(d_train, "train")]
+            eval_name = "train"
         else:
             eval_set = [(d_test, "test")]
+            eval_name = "test"
 
         steps = self.conf_params_xgboost.params.pop("steps", 300)
 
@@ -113,7 +115,7 @@ def fit(
                 num_boost_round=steps,
                 evals=eval_set,
                 verbose_eval=self.conf_xgboost.verbosity_during_final_model_training,
-                callbacks=self.get_early_stopping_callback(),
+                callbacks=self.get_early_stopping_callback(data_name=eval_name),
             )
         elif self.conf_xgboost:
             self.model = xgb.train(
@@ -123,7 +125,7 @@ def fit(
                 early_stopping_rounds=self.conf_training.early_stopping_rounds,
                 evals=eval_set,
                 verbose_eval=self.conf_xgboost.verbosity_during_final_model_training,
-                callbacks=self.get_early_stopping_callback(),
+                callbacks=self.get_early_stopping_callback(data_name=eval_name),
             )
         logging.info("Finished training")
         return self.model

bluecast/ml_modelling/xgboost_regression.py

@@ -112,8 +112,10 @@ def fit(
         d_train, d_test = self._create_d_matrices(x_train, y_train, x_test, y_test)
         if x_test.empty:
             eval_set = [(d_train, "train")]
+            eval_name = "train"
         else:
             eval_set = [(d_test, "test")]
+            eval_name = "test"
 
         steps = self.conf_params_xgboost.params.pop("steps", 300)
 
@@ -125,7 +127,7 @@ def fit(
                 early_stopping_rounds=self.conf_training.early_stopping_rounds,
                 evals=eval_set,
                 verbose_eval=self.conf_xgboost.verbosity_during_final_model_training,
-                callbacks=self.get_early_stopping_callback(),
+                callbacks=self.get_early_stopping_callback(data_name=eval_name),
             )
         elif self.conf_xgboost:
             self.model = xgb.train(
@@ -135,7 +137,7 @@ def fit(
                 early_stopping_rounds=self.conf_training.early_stopping_rounds,
                 evals=eval_set,
                 verbose_eval=self.conf_xgboost.verbosity_during_final_model_training,
-                callbacks=self.get_early_stopping_callback(),
+                callbacks=self.get_early_stopping_callback(data_name=eval_name),
             )
         logging.info("Finished training")
         return self.model

bluecast/tests/conftest.py

@@ -9,11 +9,7 @@
 
 from bluecast.blueprints.cast import BlueCast
 from bluecast.blueprints.cast_regression import BlueCastRegression
-from bluecast.config.training_config import (
-    CatboostTuneParamsConfig,
-    CatboostTuneParamsRegressionConfig,
-    TrainingConfig,
-)
+from bluecast.config.training_config import TrainingConfig
 from bluecast.tests.make_data.create_data import (
     create_synthetic_dataframe,
     create_synthetic_dataframe_regression,

bluecast/tests/test_conformal_prediction_evaluation.py

@@ -16,7 +16,7 @@
 
 def test_prediction_set_coverage():
     X, y = make_classification(
-        n_samples=1000, n_features=5, random_state=42, n_classes=2
+        n_samples=3000, n_features=5, random_state=42, n_classes=2
     )
     X_train, X_calibrate, y_train, y_calibrate = train_test_split(
         X, y, test_size=0.2, random_state=42
@@ -36,10 +36,8 @@ def test_prediction_set_coverage():
 
     # Create a custom training config and adjust general training parameters
     train_config = TrainingConfig()
-    train_config.hyperparameter_tuning_rounds = 10
-    train_config.autotune_model = (
-        False  # we want to run just normal training, no hyperparameter tuning
-    )
+    train_config.hyperparameter_tuning_rounds = 2
+    train_config.autotune_model = False
 
     automl = BlueCast(
         class_problem="binary",

bluecast/tests/test_error_paths.py

@@ -9,9 +9,7 @@
 import pytest
 
 from bluecast.blueprints.cast import BlueCast
-from bluecast.blueprints.cast_cv import BlueCastCV
 from bluecast.blueprints.cast_regression import BlueCastRegression
-from bluecast.config.training_config import TrainingConfig
 
 
 class TestPredictBeforeFit:

bluecast/tests/test_preprocessing_recipes.py

@@ -45,9 +45,9 @@ def test_initialization(preprocessing_instance):
 def test_fit_transform(sample_data, preprocessing_instance, monkeypatch):
     df, target = sample_data
 
-    # Mock the remove_correlated_columns function
+    # Mock the remove_correlated_columns function at the usage site
     monkeypatch.setattr(
-        "bluecast.preprocessing.remove_collinearity.remove_correlated_columns",
+        "bluecast.blueprints.preprocessing_recipes.remove_correlated_columns",
         mock_remove_correlated_columns,
     )
 
@@ -71,7 +71,7 @@ def test_transform(sample_data, preprocessing_instance, monkeypatch):
 
     # Fit-transform first to simulate the normal flow
     monkeypatch.setattr(
-        "bluecast.preprocessing.remove_collinearity.remove_correlated_columns",
+        "bluecast.blueprints.preprocessing_recipes.remove_correlated_columns",
         mock_remove_correlated_columns,
     )
     preprocessing_instance.fit_transform(df, target)

bluecast/tests/test_remove_collinearity.py

@@ -3,71 +3,73 @@
 from bluecast.preprocessing.remove_collinearity import remove_correlated_columns
 
 
-def test_remove_correlated_columns_high_correlation():
-    # Create a DataFrame with high correlation between columns
+def test_remove_correlated_columns_positive_correlation():
     data = {
         "A": [1, 2, 3, 4, 5],
-        "B": [2, 4, 6, 8, 10],  # B is perfectly correlated with A
-        "C": [5, 4, 3, 2, 1],  # C is not correlated with A or B
+        "B": [2, 4, 6, 8, 10],  # B is perfectly positively correlated with A
+        "C": [1, 3, 2, 5, 4],  # C has low correlation with A
     }
     df = pd.DataFrame(data)
 
     result_df = remove_correlated_columns(df, threshold=0.9)
 
-    # B should be removed because it's highly correlated with A
-    expected_df = pd.DataFrame({"A": [1, 2, 3, 4, 5], "C": [5, 4, 3, 2, 1]})
+    assert "A" in result_df.columns
+    assert "B" not in result_df.columns
+    assert "C" in result_df.columns
 
-    pd.testing.assert_frame_equal(result_df, expected_df)
 
-
-def test_remove_correlated_columns_no_removal():
-    # Create a DataFrame with no high correlations
+def test_remove_correlated_columns_negative_correlation():
     data = {
         "A": [1, 2, 3, 4, 5],
-        "B": [2, 3, 4, 5, 6],  # B is not perfectly correlated with A
-        "C": [5, 4, 3, 2, 1],  # C is not correlated with A or B
+        "B": [5, 4, 3, 2, 1],  # B is perfectly negatively correlated with A
+        "C": [1, 3, 2, 5, 4],  # C has low correlation
     }
     df = pd.DataFrame(data)
 
     result_df = remove_correlated_columns(df, threshold=0.9)
 
-    # No columns should be removed
-    pd.testing.assert_frame_equal(result_df, df)
+    assert "A" in result_df.columns
+    assert "B" not in result_df.columns, "Negative correlation should also be caught"
+    assert "C" in result_df.columns
 
 
-def test_remove_correlated_columns_no_correlation():
-    # Create a DataFrame where no columns are correlated above the threshold
+def test_remove_correlated_columns_no_removal():
     data = {
         "A": [1, 2, 3, 4, 5],
-        "B": [2, 3, 4, 5, 6],
-        "C": [5, 4, 3, 2, 1],
-        "D": [1, 2, 1, 2, 1],
+        "B": [1, 3, 2, 5, 4],  # Low correlation with A
+        "C": [3, 1, 4, 2, 5],  # Low correlation with A and B
     }
     df = pd.DataFrame(data)
 
     result_df = remove_correlated_columns(df, threshold=0.9)
 
-    # Since no columns are correlated above the threshold, the original DataFrame should be returned
-    pd.testing.assert_frame_equal(result_df, df)
+    assert list(result_df.columns) == ["A", "B", "C"]
 
 
-def test_remove_correlated_columns_different_threshold():
-    # Create a DataFrame with some correlation
+def test_remove_correlated_columns_does_not_mutate_input():
     data = {
         "A": [1, 2, 3, 4, 5],
-        "B": [2, 4, 6, 8, 10],  # B is perfectly correlated with A
-        "C": [5, 5, 5, 5, 5],  # C is constant, should have no correlation
+        "B": [2, 4, 6, 8, 10],
+        "C": [1, 3, 2, 5, 4],
     }
     df = pd.DataFrame(data)
+    original_cols = list(df.columns)
+
+    remove_correlated_columns(df, threshold=0.9)
 
-    # Use a higher threshold, so no columns should be removed
-    result_df = remove_correlated_columns(df, threshold=0.95)
+    assert list(df.columns) == original_cols, "Original DataFrame should not be mutated"
 
-    pd.testing.assert_frame_equal(result_df, df)
 
-    # Use a lower threshold, so column B should be removed
-    result_df = remove_correlated_columns(df, threshold=0.8)
+def test_remove_correlated_columns_different_threshold():
+    data = {
+        "A": [1, 2, 3, 4, 5],
+        "B": [2, 4, 6, 8, 10],  # Perfectly correlated with A
+        "C": [1, 3, 2, 5, 4],  # Low correlation
+    }
+    df = pd.DataFrame(data)
 
-    expected_df = pd.DataFrame({"A": [1, 2, 3, 4, 5], "C": [5, 5, 5, 5, 5]})
+    result_high = remove_correlated_columns(df, threshold=1.01)
+    assert len(result_high.columns) == 3, "No columns removed at threshold > 1.0"
 
-    pd.testing.assert_frame_equal(result_df, expected_df)
+    result_low = remove_correlated_columns(df, threshold=0.8)
+    assert "B" not in result_low.columns

bluecast/tests/test_save_load_roundtrip.py

@@ -4,7 +4,6 @@
 import tempfile
 
 import numpy as np
-import pytest
 
 from bluecast.general_utils.general_utils import load_for_production, save_to_production