fix: validate taxon_by column in _prep_samples_for_cohort_grouping

Signed-off-by: Suhrid Marwah <suhridmarwah07@gmail.com>

Author: suhr25

malariagen_data/anoph/frq_base.py

@@ -29,6 +29,13 @@ def _prep_samples_for_cohort_grouping(
         # Users can explicitly override with True/False.
         filter_unassigned = taxon_by == "taxon"
 
+    # Validate taxon_by.
+    if taxon_by not in df_samples.columns:
+        raise ValueError(
+            f"Invalid value for `taxon_by`: {taxon_by!r}. "
+            f"Must be the name of an existing column in the sample metadata."
+        )
+
     if filter_unassigned:
         # Remove samples with "intermediate" or "unassigned" taxon values,
         # as we only want cohorts with clean taxon calls.

tests/anoph/test_frq_base.py

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