refactoring & update unit test

src/hct_mis_api/apps/registration_datahub/mutations.py

@@ -4,7 +4,6 @@
 
 from django.core.exceptions import ValidationError
 from django.db import transaction
-from django.db.models import Q
 from django.shortcuts import get_object_or_404
 
 import graphene
@@ -50,6 +49,7 @@
     KoboImportDataNode,
     XlsxRowErrorNode,
 )
+from hct_mis_api.apps.registration_datahub.utils import get_rdi_program_population
 from hct_mis_api.apps.utils.elasticsearch_utils import (
     remove_elasticsearch_documents_by_matching_ids,
 )
@@ -115,25 +115,10 @@ def create_registration_data_import_for_import_program_population(
     business_area = BusinessArea.objects.get(slug=registration_data_import_data.pop("business_area_slug"))
     pull_pictures = registration_data_import_data.pop("pull_pictures", True)
     screen_beneficiary = registration_data_import_data.pop("screen_beneficiary", False)
-    import_from_program_id = registration_data_import_data.pop("import_from_program_id", None)
-    import_from_ids = registration_data_import_data.get("import_from_ids", "")
-
-    list_of_ids = [item.strip() for item in import_from_ids.split(",")]
-    program = get_object_or_404(Program, pk=import_to_program_id)
-    individual_ids_q = Q(id__in=list_of_ids) if program.is_social_worker_program else Q()
-    household_ids_q = Q() if program.is_social_worker_program else Q(id__in=list_of_ids)
-
-    households = Household.objects.filter(
-        household_ids_q,
-        program_id=import_from_program_id,
-        withdrawn=False,
-    ).exclude(household_collection__households__program=import_to_program_id)
-    individuals = Individual.objects.filter(
-        individual_ids_q,
-        program_id=import_from_program_id,
-        withdrawn=False,
-        duplicate=False,
-    ).exclude(individual_collection__individuals__program=import_to_program_id)
+    import_from_program_id = registration_data_import_data.pop("import_from_program_id")
+    import_from_ids = registration_data_import_data.get("import_from_ids")
+
+    households, individuals = get_rdi_program_population(import_from_program_id, import_to_program_id, import_from_ids)
     created_obj_hct = RegistrationDataImport(
         status=RegistrationDataImport.IMPORTING,
         imported_by=user,

src/hct_mis_api/apps/registration_datahub/tasks/import_program_population.py

@@ -1,38 +1,19 @@
-from django.db.models import Q
 from django.shortcuts import get_object_or_404
 
-from hct_mis_api.apps.household.models import Household, Individual
 from hct_mis_api.apps.program.models import Program
 from hct_mis_api.apps.program.utils import CopyProgramPopulation
 from hct_mis_api.apps.registration_data.models import RegistrationDataImport
+from hct_mis_api.apps.registration_datahub.utils import get_rdi_program_population
 from hct_mis_api.apps.utils.models import MergeStatusModel
 
 
 def import_program_population(
     import_from_program_id: str, import_to_program_id: str, rdi: RegistrationDataImport
 ) -> None:
-    # filter by IDs based on program DCT
-    import_from_ids = rdi.import_from_ids or ""
-    list_of_ids = [item.strip() for item in import_from_ids.split(",")]
-    import_to_program = get_object_or_404(Program, pk=import_to_program_id)
-    individual_ids_q = Q(id__in=list_of_ids) if import_to_program.is_social_worker_program else Q()
-    household_ids_q = Q() if import_to_program.is_social_worker_program else Q(id__in=list_of_ids)
-
-    copy_from_households = Household.objects.filter(
-        household_ids_q,
-        program=import_from_program_id,
-        withdrawn=False,
-    ).exclude(household_collection__households__program_id=import_to_program_id)
-    copy_from_individuals = (
-        Individual.objects.filter(
-            individual_ids_q,
-            program_id=import_from_program_id,
-            withdrawn=False,
-            duplicate=False,
-        )
-        .exclude(individual_collection__individuals__program_id=import_to_program_id)
-        .order_by("first_registration_date")
+    copy_from_households, copy_from_individuals = get_rdi_program_population(
+        import_from_program_id, import_to_program_id, rdi.import_from_ids
     )
+    import_to_program = get_object_or_404(Program, pk=import_to_program_id)
 
     CopyProgramPopulation(
         copy_from_individuals=copy_from_individuals,

src/hct_mis_api/apps/registration_datahub/utils.py

@@ -2,9 +2,14 @@
 import json
 import re
 import sys
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
+
+from django.db.models import Q, QuerySet
+from django.shortcuts import get_object_or_404
 
 from hct_mis_api.apps.core.kobo.common import get_field_name
+from hct_mis_api.apps.household.models import Household, Individual
+from hct_mis_api.apps.program.models import Program
 
 
 def post_process_dedupe_results(record: Any) -> None:
@@ -73,3 +78,48 @@ def calculate_hash_for_kobo_submission(submission: Dict) -> str:
     hash_object = hashlib.sha256(d_bytes)
     hex_dig = hash_object.hexdigest()
     return hex_dig
+
+
+def get_rdi_program_population(
+    import_from_program_id: str, import_to_program_id: str, import_from_ids: Optional[str]
+) -> Tuple[QuerySet[Household], QuerySet[Individual]]:
+    program = get_object_or_404(Program, pk=import_to_program_id)
+
+    # filter by rdi.import_from_ids HH or Ins ids based on Program.DCT
+    list_of_ids = [item.strip() for item in import_from_ids.split(",")] if import_from_ids else []
+    if list_of_ids:
+        individual_ids_q = (
+            Q(unicef_id__in=list_of_ids)
+            if program.is_social_worker_program
+            else Q(household__unicef_id__in=list_of_ids)
+        )
+        household_ids_q = (
+            Q(unicef_id__in=list_of_ids)
+            if not program.is_social_worker_program
+            else Q(individuals__unicef_id__in=list_of_ids)
+        )
+    else:
+        individual_ids_q = Q()
+        household_ids_q = Q()
+
+    households = (
+        Household.objects.filter(
+            household_ids_q,
+            program_id=import_from_program_id,
+            withdrawn=False,
+        )
+        .exclude(household_collection__households__program=import_to_program_id)
+        .distinct()
+    )
+    individuals = (
+        Individual.objects.filter(
+            individual_ids_q,
+            program_id=import_from_program_id,
+            withdrawn=False,
+            duplicate=False,
+        )
+        .exclude(individual_collection__individuals__program=import_to_program_id)
+        .distinct()
+        .order_by("first_registration_date")
+    )
+    return households, individuals

tests/unit/apps/registration_datahub/snapshots/snap_test_registration_data_import_program_population_mutation.py

@@ -110,3 +110,27 @@
         }
     ]
 }
+
+snapshots['TestRegistrationDataProgramPopulationImportMutations::test_registration_data_import_create_program_with_ids_filter 1'] = {
+    'data': {
+        'registrationProgramPopulationImport': {
+            'registrationDataImport': {
+                'dataSource': 'PROGRAM_POPULATION',
+                'name': 'New Import of Data HH Ids',
+                'screenBeneficiary': False
+            }
+        }
+    }
+}
+
+snapshots['TestRegistrationDataProgramPopulationImportMutations::test_registration_data_import_create_program_with_ids_filter 2'] = {
+    'data': {
+        'registrationProgramPopulationImport': {
+            'registrationDataImport': {
+                'dataSource': 'PROGRAM_POPULATION',
+                'name': 'New Import of Data Ind ids',
+                'screenBeneficiary': False
+            }
+        }
+    }
+}

tests/unit/apps/registration_datahub/test_registration_data_import_program_population_mutation.py

@@ -6,6 +6,7 @@
 from hct_mis_api.apps.account.permissions import Permissions
 from hct_mis_api.apps.core.base_test_case import APITestCase
 from hct_mis_api.apps.core.fixtures import create_afghanistan
+from hct_mis_api.apps.core.models import DataCollectingType
 from hct_mis_api.apps.household.fixtures import create_household_and_individuals
 from hct_mis_api.apps.program.fixtures import ProgramFactory
 from hct_mis_api.apps.program.models import Program
@@ -168,3 +169,81 @@ def test_registration_data_import_create_cannot_check_against_sanction_list(
                 }
             },
         )
+
+    def test_registration_data_import_create_program_with_ids_filter(self) -> None:
+        user = UserFactory(partner=self.partner)
+        self.create_user_role_with_permissions(user, [Permissions.RDI_IMPORT_DATA], self.afghanistan)
+        self.update_partner_access_to_program(self.partner, self.import_to_program)
+        self.household_1, self.individuals_1 = create_household_and_individuals(
+            household_data={"program": self.import_from_program},
+            individuals_data=[{}],
+        )
+        self.household_2, self.individuals_2 = create_household_and_individuals(
+            household_data={"program": self.import_from_program},
+            individuals_data=[{}],
+        )
+        self.household_3, self.individuals_3 = create_household_and_individuals(
+            household_data={"program": self.import_from_program},
+            individuals_data=[{}, {}],
+        )
+        self.household_4, self.individuals_4 = create_household_and_individuals(
+            household_data={"program": self.import_from_program},
+            individuals_data=[{}, {}],
+        )
+
+        hh_ids = f"{self.household_3.unicef_id}, {self.household_4.unicef_id}, HH-000001"
+        ind_ids = f"{self.individuals_1[0].unicef_id}, {self.individuals_2[0].unicef_id}, IND-111222"
+
+        self.snapshot_graphql_request(
+            request_string=self.CREATE_REGISTRATION_DATA_IMPORT,
+            context={
+                "user": user,
+                "headers": {"Program": self.id_to_base64(self.import_to_program.id, "ProgramNode")},
+            },
+            variables={
+                "registrationDataImportData": {
+                    "importFromProgramId": self.id_to_base64(self.import_from_program.id, "ProgramNode"),
+                    "name": "New Import of Data HH Ids",
+                    "businessAreaSlug": self.afghanistan.slug,
+                    "importFromIds": hh_ids,
+                }
+            },
+        )
+
+        rdi = RegistrationDataImport.objects.filter(name="New Import of Data HH Ids").first()
+        self.assertEqual(rdi.status, RegistrationDataImport.IMPORT_SCHEDULED)
+        self.assertEqual(rdi.data_source, RegistrationDataImport.PROGRAM_POPULATION)
+        self.assertEqual(rdi.import_from_ids, hh_ids)
+        self.assertEqual(rdi.program.data_collecting_type.type, DataCollectingType.Type.STANDARD)
+        self.assertEqual(rdi.number_of_individuals, 4)
+        self.assertEqual(rdi.number_of_households, 2)
+        self.assertEqual(rdi.program_id, self.import_to_program.id)
+
+        # update program DCT
+        self.import_to_program.data_collecting_type.type = DataCollectingType.Type.SOCIAL
+        self.import_to_program.data_collecting_type.save()
+
+        self.snapshot_graphql_request(
+            request_string=self.CREATE_REGISTRATION_DATA_IMPORT,
+            context={
+                "user": user,
+                "headers": {"Program": self.id_to_base64(self.import_to_program.id, "ProgramNode")},
+            },
+            variables={
+                "registrationDataImportData": {
+                    "importFromProgramId": self.id_to_base64(self.import_from_program.id, "ProgramNode"),
+                    "name": "New Import of Data Ind ids",
+                    "businessAreaSlug": self.afghanistan.slug,
+                    "importFromIds": ind_ids,
+                }
+            },
+        )
+
+        rdi = RegistrationDataImport.objects.filter(name="New Import of Data Ind ids").first()
+        self.assertEqual(rdi.status, RegistrationDataImport.IMPORT_SCHEDULED)
+        self.assertEqual(rdi.data_source, RegistrationDataImport.PROGRAM_POPULATION)
+        self.assertEqual(rdi.program.data_collecting_type.type, DataCollectingType.Type.SOCIAL)
+        self.assertEqual(rdi.import_from_ids, ind_ids)
+        self.assertEqual(rdi.number_of_individuals, 2)
+        self.assertEqual(rdi.number_of_households, 2)
+        self.assertEqual(rdi.program_id, self.import_to_program.id)