Merge branch 'main' into HEA-789/Create-concept-that-splits-one-expenditure-item-between-survival-or-LHP-and-Other

Author: girumb

apps/baseline/migrations/0020_alter_wealthgroup_wealth_group_category.py

@@ -0,0 +1,27 @@
+# Generated by Django 5.2.7 on 2025-11-05 10:09
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("baseline", "0019_alter_foodpurchase_times_per_year_and_more"),
+        ("metadata", "0012_alter_activitylabel_activity_type_and_more"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="wealthgroup",
+            name="wealth_group_category",
+            field=models.ForeignKey(
+                db_column="wealth_group_category_code",
+                help_text="Wealth Group Category, e.g. Poor or Better Off",
+                on_delete=django.db.models.deletion.CASCADE,
+                related_name="wealth_groups",
+                to="metadata.wealthgroupcategory",
+                verbose_name="Wealth Group Category",
+            ),
+        ),
+    ]

apps/baseline/models.py

@@ -505,6 +505,7 @@ class WealthGroup(common_models.Model):
     )
     wealth_group_category = models.ForeignKey(
         WealthGroupCategory,
+        related_name="wealth_groups",
         db_column="wealth_group_category_code",
         on_delete=models.CASCADE,
         verbose_name=_("Wealth Group Category"),

apps/baseline/tests/test_viewsets.py

@@ -1413,6 +1413,45 @@ def test_filter_by_country(self):
         self.assertEqual(response.status_code, 200)
         self.assertEqual(len(json.loads(response.content.decode("utf-8"))), 1)
 
+    def test_filter_by_product(self):
+        parent = ClassifiedProductFactory(cpc="K011")
+        product = ClassifiedProductFactory(
+            cpc="K0111",
+            description_en="my product",
+            common_name_en="common",
+            kcals_per_unit=550,
+            parent=parent,
+            aliases=["test alias"],
+        )
+        ClassifiedProductFactory(cpc="K01111")
+        characteristic1 = WealthCharacteristicFactory(
+            code="IOO", description_en="item one ownership", has_product=True
+        )
+        WealthGroupCharacteristicValueFactory(wealth_characteristic=characteristic1, product=product)
+        response = self.client.get(self.url, {"product": "K011"})
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(len(json.loads(response.content)), 1)
+        # filter by cpc
+        response = self.client.get(self.url, {"product": "K0111"})
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(len(json.loads(response.content)), 1)
+        # filter by cpc startswith
+        response = self.client.get(self.url, {"product": "K01111"})
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(len(json.loads(response.content)), 0)
+        # filter by description icontains
+        response = self.client.get(self.url, {"product": "my"})
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(len(json.loads(response.content.decode("utf-8"))), 1)
+        # filter by description
+        response = self.client.get(self.url, {"product": "my product"})
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(len(json.loads(response.content.decode("utf-8"))), 1)
+        # filter by alias
+        response = self.client.get(self.url, {"product": "test"})
+        self.assertEqual(response.status_code, 200)
+        self.assertEqual(len(json.loads(response.content.decode("utf-8"))), 1)
+
 
 class LivelihoodStrategyViewSetTestCase(APITestCase):
     @classmethod

apps/baseline/viewsets.py

@@ -471,6 +471,15 @@ class Meta:
         lookup_expr="iexact",
         label="Country",
     )
+    product = MultiFieldFilter(
+        [
+            *[(field, "icontains") for field in translation_fields("product__common_name")],
+            ("product__cpc", "istartswith"),
+            *[(field, "icontains") for field in translation_fields("product__description")],
+            ("product__aliases", "icontains"),
+        ],
+        label="Product",
+    )
 
 
 class WealthGroupCharacteristicValueViewSet(BaseModelViewSet):

apps/common/lookups.py

@@ -307,7 +307,8 @@ def do_lookup(
         if lookup_column:
             merge_df = merge_df.drop(["lookup_candidate", "lookup_key"], axis="columns")
 
-        return merge_df
+        # Preserve the original index
+        return merge_df.set_index(df.index)
 
     def get_instances(self, df, column, related_models=None):
         """

apps/common/tests/test_viewsets.py

@@ -1,3 +1,4 @@
+import importlib
 import json
 
 from rest_framework.reverse import reverse
@@ -24,6 +25,12 @@ def setUpTestData(cls):
         cls.country3 = CountryFactory()
         cls.country4 = CountryFactory()
 
+        # import baseline factory to avoid circular depdnecies
+        module = importlib.import_module("baseline.tests.factories")
+        WealthGroupFactory = getattr(module, "WealthGroupFactory")
+
+        cls.wealth_group1 = WealthGroupFactory(livelihood_zone_baseline__livelihood_zone__country=cls.country1)
+
     def setUp(self):
         self.url = reverse("country-list")
 
@@ -81,6 +88,23 @@ def test_search_by_country_name(self):
         result = json.loads(response.content.decode("utf-8"))
         self.assertEqual(len(result), 1)
 
+    def test_filter_by_has_wealthgroups(self):
+        # test by has_wealthgroups set to true
+        filter_data = {"has_wealthgroups": True}
+        response = self.client.get(self.url, filter_data)
+        self.assertEqual(response.status_code, 200)
+        result = json.loads(response.content.decode("utf-8"))
+        self.assertEqual(len(result), 1)
+        self.assertEqual(self.country1.name, result[0]["name"])
+        self.assertEqual(self.wealth_group1.livelihood_zone_baseline.livelihood_zone.country, self.country1)
+
+        # test by has_wealthgroups set to false
+        filter_data = {"has_wealthgroups": False}
+        response = self.client.get(self.url, filter_data)
+        self.assertEqual(response.status_code, 200)
+        result = json.loads(response.content.decode("utf-8"))
+        self.assertNotIn(self.country1.iso3166a2, result)
+
 
 class CurrencyViewSetTestCase(APITestCase):
     @classmethod
@@ -152,6 +176,35 @@ def setUpTestData(cls):
         cls.product2 = ClassifiedProductFactory()
         cls.superuser = UserFactory(is_superuser=True, is_staff=True, is_active=True)
 
+        cls.country_a = CountryFactory(
+            iso3166a2="AA",
+            iso3166a3="AAA",
+            iso3166n3=911,
+            iso_en_ro_name="A Country",
+            iso_en_name="AA Country",
+            name="AA Country",
+        )
+        cls.country_b = CountryFactory(
+            iso3166a2="BB",
+            iso3166a3="BBB",
+            iso3166n3=912,
+            iso_en_ro_name="B Country",
+            iso_en_name="BB Country",
+            name="BB Country",
+        )
+
+        # import baseline factory to avoid circular depdnecies
+        module = importlib.import_module("baseline.tests.factories")
+        WealthGroupFactory = getattr(module, "WealthGroupFactory")
+        LivelihoodStrategyFactory = getattr(module, "LivelihoodStrategyFactory")
+        LivelihoodZoneBaselineFactory = getattr(module, "LivelihoodZoneBaselineFactory")
+
+        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
+        )
+
     def setUp(self):
         self.url = reverse("classifiedproduct-list")
 
@@ -193,6 +246,45 @@ def test_search_fields(self):
         result = json.loads(response.content.decode("utf-8"))
         self.assertEqual(len(result), 1)
 
+    def test_filter_by_has_wealthgroups(self):
+        # test by has_wealthgroups set to true
+        filter_data = {"has_wealthgroups": True}
+        response = self.client.get(self.url, filter_data)
+        self.assertEqual(response.status_code, 200)
+        result = json.loads(response.content.decode("utf-8"))
+        self.assertEqual(len(result), 1)
+        self.assertEqual(self.product1.cpc, result[0]["cpc"])
+        self.assertEqual(self.livelihood_strategy1.product, self.product1)
+
+        # test by has_wealthgroups set to false
+        filter_data = {"has_wealthgroups": False}
+        response = self.client.get(self.url, filter_data)
+        self.assertEqual(response.status_code, 200)
+        result = json.loads(response.content.decode("utf-8"))
+        self.assertNotIn(self.product1.cpc, result)
+
+    def test_filter_by_country(self):
+        # test filter by country
+        filter_data = {"country": self.country_a.iso3166a2}
+        response = self.client.get(self.url, filter_data)
+        self.assertEqual(response.status_code, 200)
+        result = json.loads(response.content.decode("utf-8"))
+        self.assertEqual(len(result), 1)
+        self.assertEqual(self.product1.cpc, result[0]["cpc"])
+
+        filter_data = {"country": self.country_b.iso3166a2}
+        response = self.client.get(self.url, filter_data)
+        self.assertEqual(response.status_code, 200)
+        result = json.loads(response.content.decode("utf-8"))
+        self.assertEqual(len(result), 0)
+
+        filter_data = {"country": self.country_a.iso3166a2.lower()}
+        response = self.client.get(self.url, filter_data)
+        self.assertEqual(response.status_code, 200)
+        result = json.loads(response.content.decode("utf-8"))
+        self.assertEqual(len(result), 1)
+        self.assertEqual(self.product1.cpc, result[0]["cpc"])
+
 
 class UserViewSetTestCase(APITestCase):
     def setUp(self):

apps/common/viewsets.py

@@ -1,4 +1,6 @@
+from django.apps import apps
 from django.contrib.auth.models import User
+from django.db.models import Exists, OuterRef, Q
 from django.utils.text import format_lazy
 from django.utils.translation import gettext_lazy as _
 from django_filters import rest_framework as filters
@@ -8,6 +10,8 @@
 from rest_framework.pagination import PageNumberPagination
 from rest_framework.permissions import BasePermission, IsAuthenticated
 
+from common.filters import CaseInsensitiveModelMultipleChoiceFilter
+
 from .fields import translation_fields
 from .filters import MultiFieldFilter
 from .models import ClassifiedProduct, Country, Currency, UnitOfMeasure, UserProfile
@@ -98,6 +102,20 @@ def __init__(self, *args, **kwargs):
         label="Country",
     )
 
+    has_wealthgroups = filters.BooleanFilter(method="filter_has_wealthgroups")
+
+    def filter_has_wealthgroups(self, queryset, name, value):
+        if value is None:
+            return queryset
+        WealthGroup = apps.get_model("baseline", "WealthGroup")
+        wealth_group_exists = WealthGroup.objects.filter(
+            livelihood_zone_baseline__livelihood_zone__country=OuterRef("pk")
+        )
+        if value:
+            return queryset.filter(Exists(wealth_group_exists))
+        else:
+            return queryset.exclude(Exists(wealth_group_exists))
+
 
 class CountryViewSet(BaseModelViewSet):
     """
@@ -287,6 +305,35 @@ class ClassifiedProductFilterSet(filters.FilterSet):
         lookup_expr="icontains", label=format_lazy("{} ({})", _("Common Name"), _("Portuguese"))
     )
     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")
+
+    def filter_has_wealthgroups(self, queryset, name, value):
+        if value is None:
+            return queryset
+
+        WealthGroup = apps.get_model("baseline", "WealthGroup")
+
+        # Get baseline IDs that have wealth groups
+        baselines_with_wg = WealthGroup.objects.values_list("livelihood_zone_baseline_id", flat=True).distinct()
+
+        if value:
+            # Return products with strategies in those baselines
+            return queryset.filter(livelihood_strategies__livelihood_zone_baseline_id__in=baselines_with_wg).distinct()
+        else:
+            return queryset
+
+    def filter_by_country(self, queryset, name, value):
+        if not value:
+            return queryset
+
+        country_queries = Q()
+        for country in value:
+            country_queries |= Q(
+                livelihood_strategies__livelihood_zone_baseline__livelihood_zone__country__iso3166a2__iexact=country.iso3166a2
+            )
+
+        return queryset.filter(country_queries).distinct()
 
     class Meta:
         """