chg ! rebase

Author: vitali-yanushchyk-valor

src/country_workspace/contrib/aurora/forms.py

@@ -34,7 +34,7 @@ def __init__(self, *args: tuple, program: Program | None = None, **kwargs: dict)
         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):
+            if not (program.beneficiary_group and program.beneficiary_group.master_detail):
                 self.fields = {
                     key: value
                     for key, value in self.fields.items()

src/country_workspace/contrib/aurora/pipeline.py

@@ -4,11 +4,11 @@
 
 from country_workspace.contrib.aurora.client import AuroraClient
 from country_workspace.models import AsyncJob, Batch, Household, Individual
-from country_workspace.utils.config import BatchNameConfig, FailIfAlienConfig
-from country_workspace.utils.fields import uppercase_field_value, clean_field_names
+from country_workspace.utils.config import BatchNameConfig
+from country_workspace.utils.fields import clean_field_names
 
 
-class Config(BatchNameConfig, FailIfAlienConfig):
+class Config(BatchNameConfig):
     registration_reference_pk: str | None
     master_detail: bool
     household_column_prefix: NotRequired[str]
@@ -70,12 +70,11 @@ def create_household(batch: Batch, data: dict[str, Any], prefix: str) -> Househo
         ValueError: If multiple household entries are found in the provided data.
 
     """
-    flex_fields = _collect_by_prefix(data, prefix)
-    if len(flex_fields) > 1:
+    hh_data = _collect_by_prefix(data, prefix)
+    if len(hh_data) > 1:
         raise ValueError("Multiple households found")
-    flex_fields = next(iter(flex_fields.values()), {})
-    return batch.program.households.create(batch=batch, flex_fields=clean_field_names(flex_fields))
-    # return batch.program.households.create(batch=batch, flex_fields=flex_fields)
+    flex_fields = clean_field_names(next(iter(hh_data.values()), {}))
+    return batch.program.households.create(batch=batch, flex_fields=flex_fields)
 
 
 def create_individuals(
@@ -97,24 +96,21 @@ def create_individuals(
     household, individuals = None, []
     head_found = False
 
-    # for raw_individual in data.values():
-        # individual = clean_field_names(raw_individual)
-        # if not head_found:
     inds_data = _collect_by_prefix(data, cfg.get("individuals_column_prefix"))
 
     if inds_data and cfg["master_detail"] and (hh_prefix := cfg.get("household_column_prefix")):
         household = create_household(batch, data, hh_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)
+    for ind_data in inds_data.values():
+        flex_fields = clean_field_names(ind_data)
+        if household and (hh_label := cfg.get("household_label_column")) and not head_found:
+            head_found = _update_household_label_from_individual(household, flex_fields, hh_label)
         individuals.append(
             Individual(
                 batch=batch,
                 household_id=household.pk if household else None,
-                name=individual.get("given_name", ""),
-                flex_fields=individual,
+                name=flex_fields.get("given_name", ""),
+                flex_fields=flex_fields,
             )
         )
     return batch.program.individuals.bulk_create(individuals, batch_size=1000)
@@ -133,6 +129,9 @@ def _collect_by_prefix(data: dict[str, Any], prefix: str) -> dict[str, dict[str,
             and, for specific fields, values converted to uppercase. Returns an empty dictionary if no
             matching keys are found.
 
+    Raises:
+        ValueError: If a key with the specified prefix does not contain an underscore after the prefix.
+
     Examples:
         >>> data = {"user_0_relationship": "head", "user_0_gender": "male", "user_1_gender": "female"}
         >>> _collect_by_prefix(data, "user_")
@@ -144,25 +143,22 @@ 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)
-            result.setdefault(index, {})[field] = uppercase_field_value(field, 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
+            try:
+                index, field = stripped.split("_", 1)
+                result.setdefault(index, {})[field] = v
+            except ValueError:
+                raise ValueError(f"Field name '{k}' after removing prefix '{prefix}' must contain an underscore.")
+    return result
 
 
 def _update_household_label_from_individual(
-    household: Household, individual: Mapping[str, Any], household_label_column: str
+    household: Household, ind_data: Mapping[str, Any], household_label_column: str
 ) -> bool:
     """Update the household's name based on an individual's role and specified name field.
 
     Args:
         household (Household): The household instance to update.
-        individual (dict[str, Any]): A dictionary containing the individual's data,
+        ind_data (dict[str, Any]): A dictionary containing the individual's data,
             including relationship status and potential household name.
         household_label_column (str): The key in the individual's data that stores
             the name to assign to the household.
@@ -171,8 +167,8 @@ def _update_household_label_from_individual(
         bool: True if the household name was updated (individual is head and name provided), False otherwise.
 
     """
-    is_head = any(individual.get(k) == RELATIONSHIP_HEAD for k in individual if k.startswith(RELATIONSHIP_FIELDNAME))
-    name = individual.get(household_label_column)
+    is_head = any(ind_data.get(k) == RELATIONSHIP_HEAD for k in ind_data if k == RELATIONSHIP_FIELDNAME)
+    name = ind_data.get(household_label_column)
     if is_head and name:
         household.name = name
         household.save(update_fields=["name"])

src/country_workspace/utils/fields.py

@@ -11,6 +11,7 @@
 
 
 TO_REMOVE = "_h_c", "_h_f", "_i_c", "_i_f"
+TO_UPPERCASE = "relationship", "gender", "disability", "residence_status"
 
 
 def clean_field_name(v: str) -> str:
@@ -27,12 +28,21 @@ def clean_field_name(v: str) -> str:
 
 
 def clean_field_names(record: Record) -> Record:
-    return {clean_field_name(k): v for k, v in record.items()}
+    """Clean all field names in a record by normalizing them.
+
+    Args:
+        record (dict): A dictionary with field names as keys and their values.
+
+    Returns:
+        dict: A new dictionary with cleaned field names and original values.
+
+    """
+    return {clean_field_name(k): uppercase_field_value(k, v) for k, v in record.items()}
 
 
 def uppercase_field_value(k: str, v: Any) -> str:
     """
-    Convert the given field value to uppercase if applicable.
+    Convert the given field value to uppercase if its name starts with specific prefixes.
 
     Args:
         k (str): The name of the field.
@@ -42,5 +52,4 @@ def uppercase_field_value(k: str, v: Any) -> str:
         str: The uppercase value if applicable or the original value.
 
     """
-    to_uppercase = ("relationship", "gender", "disability", "residence_status")
-    return v.upper() if isinstance(v, str) and k in to_uppercase else v
+    return v.upper() if isinstance(v, str) and any(k.startswith(prefix) for prefix in TO_UPPERCASE) else v

src/country_workspace/workspaces/admin/program.py

@@ -311,8 +311,9 @@ def import_aurora(self, request: HttpRequest, program: "CountryProgram") -> "Imp
                 "batch_name": form.cleaned_data["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"],
-                "fail_if_alien": form.cleaned_data["fail_if_alien"],
-                "master_detail": (master_detail := getattr(program.beneficiary_group, "master_detail", False)),
+                "master_detail": (
+                    master_detail := (program.beneficiary_group.master_detail if program.beneficiary_group else False)
+                ),
                 **(
                     {
                         "household_column_prefix": form.cleaned_data.get("household_column_prefix"),

tests/contrib/aurora/stub.py

@@ -131,4 +131,16 @@
             },
         ],
     },
+    "invalid_key": {
+        "page": 1,
+        "results": [
+            {
+                "id": 9,
+                "flatten": {
+                    "individuals_wrong": "value",
+                    "household_invalid": "data",
+                },
+            },
+        ],
+    },
 }

tests/workspace/test_ws_import.py

@@ -6,7 +6,8 @@
 from constance import config
 from django.urls import reverse
 from pathlib import Path
-from webtest import Upload
+from webtest import forms, Upload
+
 
 from country_workspace.state import state
 from tests.contrib.aurora import stub
@@ -81,28 +82,10 @@ def test_import_data_rdi(force_migrated_records, app, program):
     assert head.name == "Edward Jeffrey Rogers"
 
 
-@pytest.mark.django_db(transaction=True)
-@pytest.mark.parametrize(
-    ("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
-        (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
-    ],
-)
-def test_import_data_aurora(
-    force_migrated_records: None,
-    app: "DjangoTestApp",
-    program: "CountryProgram",
-    mocked_responses: responses.RequestsMock,
-    stub_data: dict[str, Any],
-    error_expected: bool,
-    hh_count: int,
-    ind_count: int,
-    error_message: str,
-) -> None:
+@pytest.fixture
+def form_aurora(
+    app: "DjangoTestApp", program: "CountryProgram", mocked_responses: responses.RequestsMock, stub_data: dict[str, Any]
+) -> forms.Form:
     res = app.get("/").follow()
     res.forms["select-tenant"]["tenant"] = program.country_office.pk
     res.forms["select-tenant"].submit()
@@ -118,14 +101,30 @@ def test_import_data_aurora(
     res.forms["import-aurora"]["_selected_tab"] = "aurora"
     res.forms["import-aurora"]["aurora-registration"] = program.projects.registrations.first().pk
 
-    master_detail = program.beneficiary_group.master_detail
+    return res.forms["import-aurora"]
 
-    if error_expected and master_detail:
-        with pytest.raises(ValueError, match=error_message):
-            res.forms["import-aurora"].submit()
-        return
 
-    res = res.forms["import-aurora"].submit()
+@pytest.mark.django_db(transaction=True)
+@pytest.mark.parametrize(
+    ("stub_data", "hh_count", "ind_count"),
+    [
+        (stub.imported["correct"], 2, 3),
+        (stub.imported["no_individuals"], 0, 0),
+        (stub.imported["empty_household_data"], 1, 1),
+        (stub.imported["update_head_name"], 1, 1),
+    ],
+)
+def test_import_data_aurora_success(
+    force_migrated_records: None,
+    program: "CountryProgram",
+    form_aurora: forms.Form,
+    stub_data: dict[str, Any],
+    hh_count: int,
+    ind_count: int,
+) -> None:
+    form_aurora.submit()
+
+    master_detail = program.beneficiary_group.master_detail
     households = program.households.all()
     individuals = program.individuals.all()
     assert individuals.count() == ind_count
@@ -137,3 +136,32 @@ def test_import_data_aurora(
             assert {hh.heads().first().name for hh in households} == {"John", "Jane"}
         elif hh_count == 1 and stub_data == stub.imported["update_head_name"]:
             assert households.first().name == "Doe"
+
+
+@pytest.mark.django_db(transaction=True)
+@pytest.mark.parametrize(
+    ("stub_data", "hh_count", "ind_count", "error_message"),
+    [
+        (stub.imported["multiple_households"], 0, 1, "Multiple households found"),
+        (stub.imported["invalid_key"], 0, 0, r".*must contain an underscore"),
+    ],
+)
+def test_import_data_aurora_errors(
+    force_migrated_records: None,
+    program: "CountryProgram",
+    form_aurora: forms.Form,
+    stub_data: dict[str, Any],
+    hh_count: int,
+    ind_count: int,
+    error_message: str,
+) -> None:
+    master_detail = program.beneficiary_group.master_detail
+    expected_success = stub_data == stub.imported["multiple_households"] and not master_detail
+
+    if expected_success:
+        form_aurora.submit()
+        assert program.individuals.count() == ind_count
+        assert program.households.count() == hh_count
+    else:
+        with pytest.raises(ValueError, match=error_message):
+            form_aurora.submit()