Add test_allele_frequencies_advanced_with_sample_query_options

Author: leehart

tests/anoph/test_snp_frq.py

@@ -839,6 +839,7 @@ def check_snp_allele_frequencies_advanced(
     period_by="year",
     sample_sets=None,
     sample_query=None,
+    sample_query_options=None,
     min_cohort_size=None,
     nobs_mode="called",
     variant_query=None,
@@ -866,6 +867,7 @@ def check_snp_allele_frequencies_advanced(
         period_by=period_by,
         sample_sets=sample_sets,
         sample_query=sample_query,
+        sample_query_options=sample_query_options,
         min_cohort_size=min_cohort_size,
         nobs_mode=nobs_mode,
         variant_query=variant_query,
@@ -939,7 +941,11 @@ def check_snp_allele_frequencies_advanced(
     check_frequency(x)
 
     # Sanity check area values.
-    df_samples = api.sample_metadata(sample_sets=sample_sets, sample_query=sample_query)
+    df_samples = api.sample_metadata(
+        sample_sets=sample_sets,
+        sample_query=sample_query,
+        sample_query_options=sample_query_options,
+    )
     expected_area_values = np.unique(df_samples[area_by].dropna().values)
     area_values = ds["cohort_area"].values
     # N.B., some areas may not end up in final dataset if cohort
@@ -978,6 +984,7 @@ def check_snp_allele_frequencies_advanced(
             cohorts="admin1_year",
             sample_sets=sample_sets,
             sample_query=sample_query,
+            sample_query_options=sample_query_options,
             min_cohort_size=min_cohort_size,
             site_mask=site_mask,
             include_counts=True,
@@ -1023,6 +1030,7 @@ def check_aa_allele_frequencies_advanced(
     period_by="year",
     sample_sets=None,
     sample_query=None,
+    sample_query_options=None,
     min_cohort_size=None,
     nobs_mode="called",
     variant_query=None,
@@ -1047,6 +1055,7 @@ def check_aa_allele_frequencies_advanced(
         period_by=period_by,
         sample_sets=sample_sets,
         sample_query=sample_query,
+        sample_query_options=sample_query_options,
         min_cohort_size=min_cohort_size,
         nobs_mode=nobs_mode,
         variant_query=variant_query,
@@ -1111,7 +1120,11 @@ def check_aa_allele_frequencies_advanced(
     check_frequency(x)
 
     # Sanity checks for area values.
-    df_samples = api.sample_metadata(sample_sets=sample_sets, sample_query=sample_query)
+    df_samples = api.sample_metadata(
+        sample_sets=sample_sets,
+        sample_query=sample_query,
+        sample_query_options=sample_query_options,
+    )
     expected_area_values = np.unique(df_samples[area_by].dropna().values)
     area_values = ds["cohort_area"].values
     # N.B., some areas may not end up in final dataset if cohort
@@ -1150,6 +1163,7 @@ def check_aa_allele_frequencies_advanced(
             cohorts="admin1_year",
             sample_sets=sample_sets,
             sample_query=sample_query,
+            sample_query_options=sample_query_options,
             min_cohort_size=min_cohort_size,
             include_counts=True,
         )
@@ -1249,6 +1263,38 @@ def test_allele_frequencies_advanced_with_sample_query(
     )
 
 
+@parametrize_with_cases("fixture,api", cases=".")
+def test_allele_frequencies_advanced_with_sample_query_options(
+    fixture,
+    api: AnophelesSnpFrequencyAnalysis,
+):
+    all_sample_sets = api.sample_sets()["sample_set"].to_list()
+    df_samples = api.sample_metadata(sample_sets=all_sample_sets)
+    countries = df_samples["country"].unique().tolist()
+    countries_list = random.sample(countries, 2)
+    sample_query_options = {
+        "local_dict": {
+            "countries_list": countries_list,
+        }
+    }
+    sample_query = "country in @countries_list"
+
+    check_snp_allele_frequencies_advanced(
+        api=api,
+        sample_sets=all_sample_sets,
+        sample_query=sample_query,
+        sample_query_options=sample_query_options,
+        min_cohort_size=0,
+    )
+    check_aa_allele_frequencies_advanced(
+        api=api,
+        sample_sets=all_sample_sets,
+        sample_query=sample_query,
+        sample_query_options=sample_query_options,
+        min_cohort_size=0,
+    )
+
+
 @pytest.mark.parametrize("min_cohort_size", [0, 10, 100])
 @parametrize_with_cases("fixture,api", cases=".")
 def test_allele_frequencies_advanced_with_min_cohort_size(