Merge pull request #254 from American-Institutes-for-Research/HEA-786/fix_cascading_api

Update the country and product endpoint filter for the cascading of L…

Author: rhunwicks

apps/common/tests/test_viewsets.py

@@ -28,9 +28,30 @@ def setUpTestData(cls):
         # import baseline factory to avoid circular depdnecies
         module = importlib.import_module("baseline.tests.factories")
         WealthGroupFactory = getattr(module, "WealthGroupFactory")
+        LivelihoodZoneBaselineFactory = getattr(module, "LivelihoodZoneBaselineFactory")
+        LivelihoodStrategyFactory = getattr(module, "LivelihoodStrategyFactory")
 
         cls.wealth_group1 = WealthGroupFactory(livelihood_zone_baseline__livelihood_zone__country=cls.country1)
 
+        # Set up livelihood category cascade filter test data
+        metadata_module = importlib.import_module("metadata.tests.factories")
+        LivelihoodCategoryFactory = getattr(metadata_module, "LivelihoodCategoryFactory")
+        cls.category_a = LivelihoodCategoryFactory()
+        cls.category_b = LivelihoodCategoryFactory()
+        LivelihoodZoneBaselineFactory(
+            livelihood_zone__country=cls.country2,
+            main_livelihood_category=cls.category_a,
+        )
+
+        # Set up strategy_type cascade filter test data
+        cls.strategy_type = "CropProduction"
+        baseline_with_strategy = LivelihoodZoneBaselineFactory(livelihood_zone__country=cls.country3)
+        LivelihoodStrategyFactory(
+            livelihood_zone_baseline=baseline_with_strategy,
+            strategy_type=cls.strategy_type,
+            season__country=cls.country3,
+        )
+
     def setUp(self):
         self.url = reverse("country-list")
 
@@ -105,6 +126,37 @@ def test_filter_by_has_wealthgroups(self):
         result = json.loads(response.content.decode("utf-8"))
         self.assertNotIn(self.country1.iso3166a2, result)
 
+    def test_filter_by_livelihood_category(self):
+        # country2 has a baseline with category_a
+        response = self.client.get(self.url, {"livelihood_category": self.category_a.code})
+        self.assertEqual(response.status_code, 200)
+        result = json.loads(response.content.decode("utf-8"))
+        country_codes = [r["iso3166a2"] for r in result]
+        self.assertIn(self.country2.iso3166a2, country_codes)
+        self.assertNotIn(self.country4.iso3166a2, country_codes)
+
+        # category_b has no baselines
+        response = self.client.get(self.url, {"livelihood_category": self.category_b.code})
+        self.assertEqual(response.status_code, 200)
+        result = json.loads(response.content.decode("utf-8"))
+        self.assertEqual(len(result), 0)
+
+    def test_filter_by_strategy_type(self):
+        # country3 has a baseline with CropProduction strategy
+        response = self.client.get(self.url, {"strategy_type": self.strategy_type})
+        self.assertEqual(response.status_code, 200)
+        result = json.loads(response.content.decode("utf-8"))
+        country_codes = [r["iso3166a2"] for r in result]
+        self.assertIn(self.country3.iso3166a2, country_codes)
+        self.assertNotIn(self.country4.iso3166a2, country_codes)
+
+        # country4 has no strategies
+        response = self.client.get(self.url, {"strategy_type": "FoodPurchase"})
+        self.assertEqual(response.status_code, 200)
+        result = json.loads(response.content.decode("utf-8"))
+        country_codes = [r["iso3166a2"] for r in result]
+        self.assertNotIn(self.country4.iso3166a2, country_codes)
+
 
 class CurrencyViewSetTestCase(APITestCase):
     @classmethod
@@ -202,7 +254,28 @@ def setUpTestData(cls):
         livelihood_zone_baseline = LivelihoodZoneBaselineFactory(livelihood_zone__country=cls.country_a)
         WealthGroupFactory(livelihood_zone_baseline=livelihood_zone_baseline)
         cls.livelihood_strategy1 = LivelihoodStrategyFactory(
-            livelihood_zone_baseline=livelihood_zone_baseline, product=cls.product1
+            livelihood_zone_baseline=livelihood_zone_baseline, product=cls.product1, season__country=cls.country_a
+        )
+
+        # Set up livelihood category cascade filter test data
+        metadata_module = importlib.import_module("metadata.tests.factories")
+        LivelihoodCategoryFactory = getattr(metadata_module, "LivelihoodCategoryFactory")
+        cls.category_a = LivelihoodCategoryFactory()
+        cls.category_b = LivelihoodCategoryFactory()
+        baseline_category_a = LivelihoodZoneBaselineFactory(
+            livelihood_zone__country=cls.country_a, main_livelihood_category=cls.category_a
+        )
+        LivelihoodStrategyFactory(
+            livelihood_zone_baseline=baseline_category_a, product=cls.product1, season__country=cls.country_a
+        )
+
+        # Set up strategy_type cascade filter test data
+        cls.strategy_type_used = "CropProduction"
+        LivelihoodStrategyFactory(
+            livelihood_zone_baseline=livelihood_zone_baseline,
+            product=cls.product1,
+            strategy_type=cls.strategy_type_used,
+            season__country=cls.country_a,
         )
 
     def setUp(self):
@@ -285,6 +358,33 @@ def test_filter_by_country(self):
         self.assertEqual(len(result), 1)
         self.assertEqual(self.product1.cpc, result[0]["cpc"])
 
+    def test_filter_by_livelihood_category(self):
+        # product1 is in a baseline with category_a
+        response = self.client.get(self.url, {"livelihood_category": self.category_a.code})
+        self.assertEqual(response.status_code, 200)
+        result = json.loads(response.content.decode("utf-8"))
+        self.assertEqual(len(result), 1)
+        self.assertEqual(result[0]["cpc"], self.product1.cpc)
+
+        # category_b has no strategies
+        response = self.client.get(self.url, {"livelihood_category": self.category_b.code})
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(len(json.loads(response.content.decode("utf-8"))), 0)
+
+    def test_filter_by_strategy_type(self):
+        # product1 has a CropProduction strategy
+        response = self.client.get(self.url, {"strategy_type": self.strategy_type_used})
+        self.assertEqual(response.status_code, 200)
+        result = json.loads(response.content.decode("utf-8"))
+        self.assertEqual(len(result), 1)
+        self.assertEqual(result[0]["cpc"], self.product1.cpc)
+
+        # product2 has no strategies
+        response = self.client.get(self.url, {"strategy_type": "FoodPurchase"})
+        self.assertEqual(response.status_code, 200)
+        cpcs = [r["cpc"] for r in json.loads(response.content.decode("utf-8"))]
+        self.assertNotIn(self.product2.cpc, cpcs)
+
 
 class UserViewSetTestCase(APITestCase):
     def setUp(self):

apps/common/viewsets.py

@@ -95,6 +95,8 @@ def __init__(self, *args, **kwargs):
         self.filters["country_code"].extra["choices"] = [
             (c.pk, c.iso_en_name) for c in Country.objects.all().order_by("iso_en_name")
         ]
+        LivelihoodStrategy = apps.get_model("baseline", "LivelihoodStrategy")
+        self.filters["strategy_type"].extra["choices"] = LivelihoodStrategy._meta.get_field("strategy_type").choices
 
     country_code = filters.MultipleChoiceFilter(
         field_name="iso3166a2",
@@ -119,6 +121,16 @@ def __init__(self, *args, **kwargs):
     )
 
     has_wealthgroups = filters.BooleanFilter(method="filter_has_wealthgroups")
+    livelihood_category = filters.ModelMultipleChoiceFilter(
+        queryset=lambda request: apps.get_model("metadata", "LivelihoodCategory").objects.all(),
+        method="filter_by_livelihood_category",
+        label="Livelihood Category",
+    )
+    strategy_type = filters.MultipleChoiceFilter(
+        choices=[],
+        method="filter_by_strategy_type",
+        label="Strategy Type",
+    )
 
     def filter_has_wealthgroups(self, queryset, name, value):
         if value is None:
@@ -132,6 +144,18 @@ def filter_has_wealthgroups(self, queryset, name, value):
         else:
             return queryset.exclude(Exists(wealth_group_exists))
 
+    def filter_by_livelihood_category(self, queryset, name, value):
+        if not value:
+            return queryset
+        return queryset.filter(livelihoodzone__livelihoodzonebaseline__main_livelihood_category__in=value).distinct()
+
+    def filter_by_strategy_type(self, queryset, name, value):
+        if not value:
+            return queryset
+        return queryset.filter(
+            livelihoodzone__livelihoodzonebaseline__livelihood_strategies__strategy_type__in=value
+        ).distinct()
+
 
 class CountryViewSet(BaseModelViewSet):
     """
@@ -289,6 +313,11 @@ class ClassifiedProductFilterSet(filters.FilterSet):
                          The filter will display choices based on the available UnitOfMeasure objects.
     """
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        LivelihoodStrategy = apps.get_model("baseline", "LivelihoodStrategy")
+        self.filters["strategy_type"].extra["choices"] = LivelihoodStrategy._meta.get_field("strategy_type").choices
+
     cpc = filters.CharFilter(lookup_expr="icontains", label="CPC v2.1")
     description_en = filters.CharFilter(
         lookup_expr="icontains", label=format_lazy("{} ({})", _("Description"), _("English"))
@@ -323,6 +352,16 @@ class ClassifiedProductFilterSet(filters.FilterSet):
     unit_of_measure = filters.ModelChoiceFilter(queryset=UnitOfMeasure.objects.all(), field_name="unit_of_measure")
     has_wealthgroups = filters.BooleanFilter(method="filter_has_wealthgroups")
     country = CaseInsensitiveModelMultipleChoiceFilter(queryset=Country.objects.all(), method="filter_by_country")
+    livelihood_category = filters.ModelMultipleChoiceFilter(
+        queryset=lambda request: apps.get_model("metadata", "LivelihoodCategory").objects.all(),
+        method="filter_by_livelihood_category",
+        label="Livelihood Category",
+    )
+    strategy_type = filters.MultipleChoiceFilter(
+        choices=[],
+        method="filter_by_strategy_type",
+        label="Strategy Type",
+    )
 
     def filter_has_wealthgroups(self, queryset, name, value):
         if value is None:
@@ -351,6 +390,18 @@ def filter_by_country(self, queryset, name, value):
 
         return queryset.filter(country_queries).distinct()
 
+    def filter_by_livelihood_category(self, queryset, name, value):
+        if not value:
+            return queryset
+        return queryset.filter(
+            livelihood_strategies__livelihood_zone_baseline__main_livelihood_category__in=value
+        ).distinct()
+
+    def filter_by_strategy_type(self, queryset, name, value):
+        if not value:
+            return queryset
+        return queryset.filter(livelihood_strategies__strategy_type__in=value).distinct()
+
     class Meta:
         """
         Metadata options for the ClassifiedProductFilterSet.