chg ! aurora import takes into account beneficiary groups

Author: vitali-yanushchyk-valor

src/country_workspace/contrib/aurora/forms.py

@@ -6,37 +6,37 @@
 
 class ImportAuroraForm(forms.Form):
     batch_name = forms.CharField(required=False, help_text="Label for this batch.")
-
     registration = forms.ModelChoiceField(
         queryset=Registration.objects.none(),
         help_text="What type of registrations are being imported.",
     )
-
     household_column_prefix = forms.CharField(
-        initial="household_",
-        help_text="Household's column group prefix",
+        initial="household_", help_text="Household's column group prefix", required=False
     )
-
     individuals_column_prefix = forms.CharField(
         initial="individuals_",
         help_text="Individuals' column group prefix",
     )
-
     household_label_column = forms.CharField(
         required=False,
         initial="family_name",
         help_text="Which Individual's column should be used as label for the household.",
     )
-
     check_before = forms.BooleanField(
         required=False, help_text="Prevent import if errors if data is not valid against data checker."
     )
-
     fail_if_alien = forms.BooleanField(
         required=False, help_text="Fails if it finds fields which do not exists in data checker."
     )
 
     def __init__(self, *args: tuple, program: Program | None = None, **kwargs: dict) -> None:
         super().__init__(*args, **kwargs)
+        self.program = program
         if program:
             self.fields["registration"].queryset = Registration.objects.filter(project__program=program, active=True)
+            if not getattr(program.beneficiary_group, "master_detail", False):
+                self.fields = {
+                    key: value
+                    for key, value in self.fields.items()
+                    if key not in ("household_column_prefix", "household_label_column")
+                }

src/country_workspace/contrib/aurora/pipeline.py

@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Any, TypedDict, cast, Final, NotRequired
 
 from django.db.transaction import atomic
 
@@ -7,52 +7,61 @@
 from country_workspace.utils.fields import clean_field_name, uppercase_field_value
 
 
+class Config(TypedDict):
+    batch_name: str
+    registration_reference_pk: int | None
+    master_detail: bool
+    household_column_prefix: NotRequired[str]
+    individuals_column_prefix: str
+    household_label_column: NotRequired[str]
+    check_before: bool
+    fail_if_alien: bool
+
+
+HEAD_RELATIONSHIP: Final[str] = "HEAD"
+
+
 def import_from_aurora(job: AsyncJob) -> dict[str, int]:
     """Import data from the Aurora system into the database within an atomic transaction.
 
     Args:
         job (AsyncJob): The job instance containing the configuration and context for data import.
-            Expected keys in `job.config`:
-            - "batch_name" (str): The name for the newly created batch.
-            - "registration_reference_pk" (int): The unique identifier of the registration to import.
-            - "household_column_prefix" (str, optional): The prefix for household-related columns.
-            - "individuals_column_prefix" (str, optional): The prefix for individual-related columns.
-            - "household_label_column" (str, optional): The column name used to determine the household label.
+            Expected keys in `job.config` correspond to the `Config` TypedDict.
 
     Returns:
-        dict[str, int]: A dictionary with the counts of successfully created records:
-            - "households": The number of households imported.
-            - "individuals": The total number of individuals imported.
+        dict[str, int]: Counts of imported records:
+            - "households": Number of households imported (0 if `master_detail` is False or None).
+            - "individuals": Total number of individuals imported.
 
     """
-    total_hh = total_ind = 0
+    total = {"households": 0, "individuals": 0}
+    cfg = cast(Config, job.config)
+
     batch = Batch.objects.create(
-        name=job.config["batch_name"],
+        name=cfg["batch_name"],
         program=job.program,
         country_office=job.program.country_office,
         imported_by=job.owner,
         source=Batch.BatchSource.AURORA,
     )
+
     client = AuroraClient()
     with atomic():
-        for record in client.get(f"registration/{job.config['registration_reference_pk']}/records/"):
-            inds_data = _collect_by_prefix(record["flatten"], job.config.get("individuals_column_prefix"))
-            if inds_data:
-                hh = create_household(batch, record["flatten"], job.config.get("household_column_prefix"))
-                total_hh += 1
-                total_ind += len(
-                    create_individuals(
-                        household=hh,
-                        data=inds_data,
-                        household_label_column=job.config.get("household_label_column"),
-                    )
-                )
-    return {"households": total_hh, "individuals": total_ind}
+        for record in client.get(f"registration/{cfg['registration_reference_pk']}/records/"):
+            individuals = create_individuals(
+                batch=batch,
+                data=record["flatten"],
+                cfg=cfg,
+            )
+            if cfg.get("master_detail") and individuals and individuals[0].household_id:
+                total["households"] += 1
+            total["individuals"] += len(individuals)
+
+    return total
 
 
 def create_household(batch: Batch, data: dict[str, Any], prefix: str) -> Household:
-    """
-    Create a Household object from the provided data and associate it with a batch.
+    """Create a Household object from the provided data and associate it with a batch.
 
     Args:
         batch (Batch): The batch to which the household will be linked.
@@ -69,36 +78,47 @@ def create_household(batch: Batch, data: dict[str, Any], prefix: str) -> Househo
     flex_fields = _collect_by_prefix(data, prefix)
     if len(flex_fields) > 1:
         raise ValueError("Multiple households found")
+    flex_fields = next(iter(flex_fields.values()), {})
     return batch.program.households.create(batch=batch, flex_fields=flex_fields)
 
 
-def create_individuals(household: Household, data: dict[str, Any], household_label_column: str) -> list[Individual]:
-    """Create and associate Individual objects with a given Household.
+def create_individuals(
+    batch: Batch,
+    data: dict[str, Any],
+    cfg: Config,
+) -> list[Individual]:
+    """Create and associate Individual objects with an optional Household.
 
     Args:
-        household (Household): The household to which the individuals will be linked.
-        data (dict[str, Any]): A dictionary mapping indices to individual details.
-        household_label_column (str): The key in the individual data used to determine the household label.
+        batch (Batch): The batch to which individuals will be linked.
+        data (dict[str, Any]): A dictionary containing related information.
+        cfg (Config): Configuration dictionary containing various settings for the import process.
 
     Returns:
         list[Individual]: A list of successfully created Individual instances.
 
     """
-    individuals = []
+    household, individuals = None, []
     head_found = False
 
-    for individual in data.values():
-        if not head_found:
+    inds_data = _collect_by_prefix(data, cfg.get("individuals_column_prefix"))
+
+    if all((inds_data, cfg["master_detail"], cfg.get("household_column_prefix"))):
+        household = create_household(batch, data, cfg.get("household_column_prefix"))
+
+    for individual in inds_data.values():
+        household_label_column = cfg.get("household_label_column")
+        if household and household_label_column and not head_found:
             head_found = _update_household_label_from_individual(household, individual, household_label_column)
         individuals.append(
             Individual(
-                batch=household.batch,
-                household_id=household.pk,
+                batch=batch,
+                household_id=household.pk if household else None,
                 name=individual.get("given_name", ""),
                 flex_fields=individual,
-            ),
+            )
         )
-    return household.program.individuals.bulk_create(individuals)
+    return batch.program.individuals.bulk_create(individuals, batch_size=1000)
 
 
 def _collect_by_prefix(data: dict[str, Any], prefix: str) -> dict[str, dict[str, Any]]:
@@ -123,11 +143,12 @@ def _collect_by_prefix(data: dict[str, Any], prefix: str) -> dict[str, dict[str,
 
     """
     result = {}
-    for k, v in data.items():
-        if (stripped := k.removeprefix(prefix)) != k:
-            index, field = stripped.split("_", 1)
-            field_clean = clean_field_name(field)
-            result.setdefault(index, {})[field_clean] = uppercase_field_value(field_clean, v)
+    for key, value in data.items():
+        if not key.startswith(prefix):
+            continue
+        index, field = key.removeprefix(prefix).split("_", 1)
+        clean_field = clean_field_name(field)
+        result.setdefault(index, {})[clean_field] = uppercase_field_value(clean_field, value)
     return result
 
 

src/country_workspace/workspaces/admin/program.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 from admin_extra_buttons.api import button
 from django import forms
@@ -13,7 +13,7 @@
 from django.utils.translation import gettext as _
 from strategy_field.utils import fqn
 
-from country_workspace.contrib.aurora.pipeline import import_from_aurora
+from country_workspace.contrib.aurora.pipeline import import_from_aurora, Config as AuroraConfig
 from country_workspace.state import state
 
 from ...contrib.aurora.forms import ImportAuroraForm
@@ -306,21 +306,33 @@ def import_rdi(self, request: HttpRequest, program: CountryProgram) -> "ImportFi
     def import_aurora(self, request: HttpRequest, program: "CountryProgram") -> "ImportAuroraForm|None":
         form = ImportAuroraForm(request.POST, prefix="aurora", program=program)
         if form.is_valid():
-            registration_reference_pk = getattr(form.cleaned_data["registration"], "reference_pk", None)
+            config = cast(
+                AuroraConfig,
+                {
+                    "batch_name": form.cleaned_data.get("batch_name") or batch_name_default(),
+                    "registration_reference_pk": getattr(form.cleaned_data.get("registration"), "reference_pk", None),
+                    "individuals_column_prefix": form.cleaned_data["individuals_column_prefix"],
+                    "check_before": form.cleaned_data.get("check_before", False),
+                    "fail_if_alien": form.cleaned_data.get("fail_if_alien", False),
+                    "master_detail": (master_detail := getattr(program.beneficiary_group, "master_detail", False)),
+                    **(
+                        {
+                            "household_column_prefix": form.cleaned_data.get("household_column_prefix"),
+                            "household_label_column": form.cleaned_data.get("household_label_column"),
+                        }
+                        if master_detail
+                        else {}
+                    ),
+                },
+            )
             job: AsyncJob = AsyncJob.objects.create(
                 description="Aurora importing",
                 type=AsyncJob.JobType.TASK,
                 action=fqn(import_from_aurora),
                 file=None,
                 program=program,
                 owner=request.user,
-                config={
-                    "batch_name": form.cleaned_data["batch_name"] or batch_name_default(),
-                    "registration_reference_pk": registration_reference_pk,
-                    "household_column_prefix": form.cleaned_data["household_column_prefix"],
-                    "individuals_column_prefix": form.cleaned_data["individuals_column_prefix"],
-                    "household_label_column": form.cleaned_data["household_label_column"],
-                },
+                config=config,
             )
             job.queue()
             self.message_user(request, _("Import scheduled"), messages.SUCCESS)

tests/workspace/test_ws_import.py

@@ -27,8 +27,8 @@ def office():
     return co
 
 
-@pytest.fixture
-def program(office, force_migrated_records, household_checker, individual_checker):
+@pytest.fixture(params=[True, False], ids=["master_detail_true", "master_detail_false"])
+def program(request, office, force_migrated_records, household_checker, individual_checker):
     from testutils.factories import CountryProgramFactory, ProjectFactory, RegistrationFactory
 
     program = CountryProgramFactory(
@@ -37,6 +37,7 @@ def program(office, force_migrated_records, household_checker, individual_checke
         individual_checker=individual_checker,
         household_columns="name\nid\nxx",
         individual_columns="name\nid\nxx",
+        beneficiary_group__master_detail=request.param,
     )
     project = ProjectFactory(program=program)
     RegistrationFactory.create_batch(3, project=project, active=True)
@@ -85,8 +86,8 @@ def test_import_data_rdi(force_migrated_records, app, program):
     ("stub_data", "error_expected", "hh_count", "ind_count", "error_message"),
     [
         (stub.imported["correct"], False, 2, 3, None),  # 2 hh: 1st with 1 ind, 2nd with 2 inds
-        (stub.imported["no_individuals"], False, 0, 0, None),  # No individuals, no Household
-        (stub.imported["multiple_households"], True, 0, 0, "Multiple households found"),  # Multiple households error
+        (stub.imported["no_individuals"], False, 0, 0, None),  # No individuals
+        (stub.imported["multiple_households"], True, 0, 1, "Multiple households found"),  # Multiple households error
         (stub.imported["empty_household_data"], False, 1, 1, None),  # Only ind without hh data
         (stub.imported["update_head_name"], False, 1, 1, None),  # Household name updated from head
     ],
@@ -107,7 +108,6 @@ def test_import_data_aurora(
     res.forms["select-tenant"].submit()
 
     url = reverse("workspace:workspaces_countryprogram_import_data", args=[program.pk])
-
     mocked_responses.add(
         responses.GET,
         re.compile(re.escape(config.AURORA_API_URL) + ".*"),
@@ -118,14 +118,20 @@ def test_import_data_aurora(
     res.forms["import-aurora"]["_selected_tab"] = "aurora"
     res.forms["import-aurora"]["aurora-registration"] = program.projects.registrations.first().pk
 
-    if error_expected:
+    master_detail = program.beneficiary_group.master_detail
+
+    if error_expected and master_detail:
         with pytest.raises(ValueError, match=error_message):
             res.forms["import-aurora"].submit()
-    else:
-        res = res.forms["import-aurora"].submit()
-        households = program.households.all()
-        assert households.count() == hh_count
-        assert sum(hh.members.count() for hh in households) == ind_count
+        return
+
+    res = res.forms["import-aurora"].submit()
+    households = program.households.all()
+    individuals = program.individuals.all()
+    assert individuals.count() == ind_count
+    assert households.count() == (hh_count if master_detail else 0)
+
+    if master_detail:
         if hh_count == 2:
             assert {hh.members.count() for hh in households} == {1, 2}
             assert {hh.heads().first().name for hh in households} == {"John", "Jane"}