Merge pull request #1307 from BLSQ/POLIO-1548-hide-non-polio-campaign

POLIO-1548 Hide non-polio campaigns

plugins/polio/budget/api.py

@@ -78,7 +78,7 @@ def destroy(self, request, *args, **kwargs):
 
     def get_queryset(self) -> QuerySet:
         user = self.request.user
-        campaigns = Campaign.objects.filter_for_user(user)
+        campaigns = Campaign.polio_objects.filter_for_user(user)
         budget_processes = (
             BudgetProcess.objects.filter(rounds__campaign__in=campaigns)
             .distinct()
@@ -125,13 +125,19 @@ def available_rounds_for_create(self, request):
         """
         Returns all available rounds that can be used to create a new `BudgetProcess`.
         """
-        user_campaigns = Campaign.objects.filter_for_user(self.request.user).filter(country__isnull=False)
+        user_campaigns = Campaign.polio_objects.filter_for_user(self.request.user).filter(country__isnull=False)
         available_rounds = (
             Round.objects.filter(budget_process__isnull=True, campaign__in=user_campaigns)
             .select_related("campaign__country")
             .order_by("campaign__country__name", "campaign__obr_name", "number")
             .only(
-                "id", "number", "campaign_id", "campaign__obr_name", "campaign__country_id", "campaign__country__name"
+                "id",
+                "number",
+                "campaign_id",
+                "campaign__obr_name",
+                "campaign__country_id",
+                "campaign__country__name",
+                "target_population",
             )
         )
         return Response(available_rounds.as_ui_dropdown_data(), status=status.HTTP_200_OK)
@@ -146,14 +152,20 @@ def available_rounds_for_update(self, request):
         campaign_uuid = query_params.validated_data["campaign_id"]
         budget_process_id = query_params.validated_data["budget_process_id"]
 
-        campaign = Campaign.objects.filter(id=campaign_uuid).filter_for_user(self.request.user).first()
+        campaign = Campaign.polio_objects.filter(id=campaign_uuid).filter_for_user(self.request.user).first()
         available_rounds = (
             Round.objects.filter(campaign=campaign)
             .select_related("campaign__country")
             .filter(Q(budget_process_id=budget_process_id) | Q(budget_process__isnull=True))
             .order_by("number")
             .only(
-                "id", "number", "campaign_id", "campaign__obr_name", "campaign__country_id", "campaign__country__name"
+                "id",
+                "number",
+                "campaign_id",
+                "campaign__obr_name",
+                "campaign__country_id",
+                "campaign__country__name",
+                "target_population",
             )
         )
         return Response(available_rounds.as_ui_dropdown_data()["rounds"], status=status.HTTP_200_OK)

plugins/polio/models.py

@@ -321,26 +321,6 @@ def districts_count_calculated(self):
         return len(self.campaign.get_districts_for_round(self))
 
 
-class CampaignQuerySet(models.QuerySet):
-    def filter_for_user(self, user: Union[User, AnonymousUser]):
-        qs = self
-        if user.is_authenticated:
-            # Authenticated users only get campaigns linked to their account
-            qs = qs.filter(account=user.iaso_profile.account)
-
-            # Restrict Campaign to the OrgUnit on the country he can access
-            if user.iaso_profile.org_units.count():
-                org_units = OrgUnit.objects.hierarchy(user.iaso_profile.org_units.all()).defer(
-                    "geom", "simplified_geom"
-                )
-                qs = qs.filter(Q(country__in=org_units) | Q(initial_org_unit__in=org_units))
-        return qs
-
-
-# workaround for MyPy detection
-CampaignManager = models.Manager.from_queryset(CampaignQuerySet)
-
-
 class CampaignType(models.Model):
     POLIO = "Polio"
     MEASLES = "Measles"
@@ -364,11 +344,35 @@ def save(self, *args, **kwargs):
         super(CampaignType, self).save(*args, **kwargs)
 
 
+class CampaignQuerySet(models.QuerySet):
+    def filter_for_user(self, user: Union[User, AnonymousUser]):
+        qs = self
+        if user.is_authenticated:
+            # Authenticated users only get campaigns linked to their account
+            qs = qs.filter(account=user.iaso_profile.account)
+
+            # Restrict Campaign to the OrgUnit on the country he can access
+            if user.iaso_profile.org_units.count():
+                org_units = OrgUnit.objects.hierarchy(user.iaso_profile.org_units.all()).defer(
+                    "geom", "simplified_geom"
+                )
+                qs = qs.filter(Q(country__in=org_units) | Q(initial_org_unit__in=org_units))
+        return qs
+
+
+class PolioCampaignManager(models.Manager):
+    def get_queryset(self):
+        return super().get_queryset().prefetch_related("campaign_types").filter(campaign_types__name=CampaignType.POLIO)
+
+
 class Campaign(SoftDeletableModel):
     class Meta:
         ordering = ["obr_name"]
 
-    objects = CampaignManager()
+    # Managers.
+    objects = models.Manager.from_queryset(CampaignQuerySet)()
+    polio_objects = PolioCampaignManager.from_queryset(CampaignQuerySet)()
+
     scopes: "django.db.models.manager.RelatedManager[CampaignScope]"
     rounds: "django.db.models.manager.RelatedManager[Round]"
     id = models.UUIDField(default=uuid4, primary_key=True, editable=False)

plugins/polio/tests/api/test_budget.py

@@ -13,7 +13,7 @@
 from iaso import models as m
 from iaso.test import APITestCase
 from plugins.polio.budget.models import BudgetProcess, BudgetStep, MailTemplate
-from plugins.polio.models import Campaign, Round
+from plugins.polio.models import Campaign, CampaignType, Round
 from plugins.polio.tests.utils.budget import get_mocked_workflow
 
 
@@ -48,12 +48,16 @@ def setUpTestData(cls) -> None:
             username="test", first_name="test", last_name="test", account=cls.account, permissions=["iaso_polio_budget"]
         )
 
+        # Campaign type.
+        cls.polio_type = CampaignType.objects.get(name=CampaignType.POLIO)
+
         # Campaign.
         cls.campaign = Campaign.objects.create(
             obr_name="test campaign",
             account=cls.user.iaso_profile.account,
             country=m.OrgUnit.objects.create(name="ANGOLA"),
         )
+        cls.campaign.campaign_types.set([cls.polio_type])
 
         # Budget Processes.
         cls.budget_process_1 = BudgetProcess.objects.create(created_by=cls.user)
@@ -789,13 +793,17 @@ def setUpTestData(cls) -> None:
         cls.account = m.Account.objects.create(name="Account", default_version=cls.source_version)
         cls.user = cls.create_user_with_profile(username="user", account=cls.account, permissions=["iaso_polio_budget"])
 
+        # Campaign type.
+        cls.polio_type = CampaignType.objects.get(name=CampaignType.POLIO)
+
         # Campaign.
         cls.country = m.OrgUnit.objects.create(name="ANGOLA")
         cls.campaign = Campaign.objects.create(
             obr_name="test campaign",
             account=cls.user.iaso_profile.account,
             country=cls.country,
         )
+        cls.campaign.campaign_types.set([cls.polio_type])
 
         # Budget Processes.
         cls.budget_process_1 = BudgetProcess.objects.create(created_by=cls.user)

plugins/polio/tests/models/test_campaign_models.py

@@ -0,0 +1,42 @@
+from iaso import models as m
+from iaso.test import TestCase
+
+from plugins.polio.models import Campaign, CampaignType
+
+
+class CampaignTestCase(TestCase):
+    """
+    Test Campaign model.
+    """
+
+    @classmethod
+    def setUpTestData(cls):
+        # User.
+        cls.data_source = m.DataSource.objects.create(name="Data Source")
+        cls.source_version = m.SourceVersion.objects.create(data_source=cls.data_source, number=1)
+        cls.account = m.Account.objects.create(name="Account", default_version=cls.source_version)
+        cls.user = cls.create_user_with_profile(
+            email="[email protected]",
+            username="test",
+            first_name="John",
+            last_name="Doe",
+            account=cls.account,
+            permissions=["iaso_polio_budget"],
+        )
+
+        # Campaigns types.
+        cls.polio_type = CampaignType.objects.get(name=CampaignType.POLIO)
+        cls.measles_type = CampaignType.objects.get(name=CampaignType.MEASLES)
+
+        # Campaigns.
+        cls.polio_campaign = Campaign.objects.create(obr_name="Polio", account=cls.account)
+        cls.polio_campaign.campaign_types.set([cls.polio_type])
+        cls.measles_campaign = Campaign.objects.create(obr_name="Measles", account=cls.account)
+        cls.measles_campaign.campaign_types.set([cls.measles_type])
+
+    def test_polio_campaign_manager(self):
+        self.assertEqual(Campaign.objects.count(), 2)
+        self.assertEqual(Campaign.polio_objects.count(), 1)
+
+        polio_campaign = Campaign.polio_objects.first()
+        self.assertEqual(polio_campaign.campaign_types.first().name, CampaignType.POLIO)