Fix test

Author: sarakodeiri

examples/ept_attack/config.yaml

@@ -16,5 +16,8 @@ pipeline:
   run_feature_extraction: true # Whether to run attribute prediction model training
   run_attack_classifier_training: false # Whether to run attack classifier training
 
+attack_settings:
+  single_table: true # Whether the data is single-table
+
 # General settings
 random_seed: 42

src/midst_toolkit/attacks/ept/feature_extraction.py

@@ -54,7 +54,6 @@ def preprocess_train_predict(
             y_test: True values for the target column on the test data.
             task_type: Whether the attribution prediction model was a classification or regression model.
     """
-
     assert target_col in train_points.columns, f"Target column '{target_col}' not found in train_points."
     assert target_col in test_points.columns, f"Target column '{target_col}' not found in test_points."
 
@@ -162,7 +161,7 @@ def extract_features(
         columns.append(column)
 
         if task_type == TaskType.CLASSIFICATION:
-            #TODO: Maybe change the variable name from accuracy to correctness
+            # TODO: Maybe change the variable name from accuracy to correctness
             # Calculate accuracy
             accuracy = predictions == y_test
             accuracy = accuracy.astype(int)

tests/unit/attacks/ept_attack/test_feature_extraction.py

@@ -81,20 +81,33 @@ def test_preprocess_train_predict_assertions(sample_dataframes, sample_column_ty
     # Tests that the assertions within preprocess_train_predict fire correctly.
 
     train_df, test_df = sample_dataframes
+    target_col = "num_col_1"
 
-    # Test mismatching columns
-    test_df_mismatch = test_df.drop(columns=["num_col_1"])
-    with pytest.raises(AssertionError, match="Columns in df_train and df_test do not match"):
-        preprocess_train_predict(train_df, test_df_mismatch, "cat_col_1", sample_column_types)
+    # 1. Target column not in train_points
+    with pytest.raises(AssertionError, match="Target column 'non_existent_col' not found in train_points."):
+        preprocess_train_predict(train_df, test_df, "non_existent_col", sample_column_types)
 
-    # Test target_col not in column_types
-    with pytest.raises(AssertionError, match="must appear exactly once"):
-        preprocess_train_predict(train_df, test_df, "missing_col", sample_column_types)
+    # 2. Target column not in test_points
+    test_df_missing_target = test_df.drop(columns=[target_col])
+    with pytest.raises(AssertionError, match=f"Target column '{target_col}' not found in test_points."):
+        preprocess_train_predict(train_df, test_df_missing_target, target_col, sample_column_types)
 
-    # Test column_types not matching dataframe columns
-    bad_column_types = {"numerical": ["num_col_1"], "categorical": []}
-    with pytest.raises(AssertionError, match="must match the columns in the combined dataframe"):
-        preprocess_train_predict(train_df, test_df, "num_col_1", bad_column_types)
+    # 3. Mismatched columns between train and test dataframes
+    test_df_mismatched = test_df.rename(columns={"num_col_2": "new_col_name"})
+    with pytest.raises(AssertionError, match="Columns in df_train and df_test do not match"):
+        preprocess_train_predict(train_df, test_df_mismatched, target_col, sample_column_types)
+
+    # 4. Target column appears more than once in column_types
+    column_types_duplicate = sample_column_types.copy()
+    column_types_duplicate["categorical"] = column_types_duplicate["categorical"] + [target_col]
+    with pytest.raises(AssertionError, match=f"The target column '{target_col}' must appear exactly once"):
+        preprocess_train_predict(train_df, test_df, target_col, column_types_duplicate)
+
+    # 5. Mismatch between dataframe columns and column_types
+    column_types_mismatch = sample_column_types.copy()
+    column_types_mismatch["numerical"] = ["num_col_1"]  # Missing num_col_2
+    with pytest.raises(AssertionError, match="The union of numeric_columns and categorical_columns must match"):
+        preprocess_train_predict(train_df, test_df, target_col, column_types_mismatch)
 
 
 def test_main_feature_extraction(sample_dataframes, sample_column_types):