Merge pull request #1140 from suhr25/fix/area-by-validation

fix: validate area_by column in _prep_samples_for_cohort_grouping

Author: jonbrenas

malariagen_data/anoph/frq_base.py

@@ -78,6 +78,13 @@ def _prep_samples_for_cohort_grouping(
         # Apply the matching period_by function to create a new "period" column.
         df_samples["period"] = df_samples.apply(period_by_func, axis="columns")
 
+    # Validate area_by.
+    if area_by not in df_samples.columns:
+        raise ValueError(
+            f"Invalid value for `area_by`: {area_by!r}. "
+            f"Must be the name of an existing column in the sample metadata."
+        )
+
     # Copy the specified area_by column to a new "area" column.
     df_samples["area"] = df_samples[area_by]
 

tests/anoph/test_frq_base.py

@@ -101,6 +101,23 @@ def test_does_not_modify_original(self):
         assert df["taxon"].tolist() == original_values
 
 
+class TestPrepSamplesAreaByValidation:
+    """Tests for area_by validation in _prep_samples_for_cohort_grouping."""
+
+    def test_invalid_area_by_raises_value_error(self):
+        """A non-existent area_by column should raise ValueError, not KeyError."""
+        import pytest
+
+        df = _make_test_df()
+        with pytest.raises(ValueError, match="Invalid value for `area_by`"):
+            _prep_samples_for_cohort_grouping(
+                df_samples=df,
+                area_by="nonexistent_column",
+                period_by="year",
+                taxon_by="taxon",
+            )
+
+
 class TestPlotFrequenciesTimeSeriesMissingCI:
     """Tests for plot_frequencies_time_series when CI variables are absent.