Refactor batch transformer application in reprocessing.py

- Removed the `_apply_transformer_only` function to streamline the transformation process.
- Updated `apply_batch_transformers` to delegate transformation logic to `run_batch_postprocessing`, enhancing clarity and maintainability.
- Added new unit tests for `apply_batch_transformers` to ensure proper error handling and validation of transformer IDs.
- Improved test coverage for batch processing scenarios, verifying the correct delegation and response structure.

Author: Arsen P

src/country_workspace/workspaces/admin/batch/reprocessing.py

@@ -16,8 +16,7 @@
     build_household_processor as build_kobo_household_processor,
     build_individual_processor as build_kobo_individual_processor,
 )
-from country_workspace.models import AsyncJob, Batch, Household, Individual, MappingImporter, Program, Transformer
-from country_workspace.utils.collector_linkage import sync_collector_links
+from country_workspace.models import AsyncJob, Batch, Household, Individual, MappingImporter, Program
 from country_workspace.utils.fields import to_reference_key
 from country_workspace.utils.import_flow import run_batch_postprocessing, build_import_processor
 from country_workspace.workspaces.admin.cleaners.validate import create_validation_jobs
@@ -95,17 +94,6 @@ def _apply_import_processor(
     return True
 
 
-def _apply_transformer_only(record: Household | Individual, transformer: Transformer) -> bool:
-    data = record.flex_fields if isinstance(record.flex_fields, dict) else {}
-    transformed = transformer.apply(data.copy())
-    if transformed == record.flex_fields:
-        return False
-
-    record.flex_fields = transformed
-    record.save(update_fields=["flex_fields"])
-    return True
-
-
 def _build_processor(
     *,
     batch: Batch,
@@ -403,42 +391,33 @@ def reprocess_batch(job: AsyncJob) -> dict[str, Any]:
 def apply_batch_transformers(job: AsyncJob) -> dict[str, Any]:
     batch_id, batch = _get_batch_from_job(job)
 
-    household_transformer_id, household_transformer = _resolve_config_object(
+    household_transformer_id, _ = _resolve_config_object(
         batch.country_office.transformers.all(),
         job.config.get("household_transformer_id"),
         "Household transformer",
     )
-    individual_transformer_id, individual_transformer = _resolve_config_object(
+    individual_transformer_id, _ = _resolve_config_object(
         batch.country_office.transformers.all(),
         job.config.get("individual_transformer_id"),
         "Individual transformer",
     )
     if not household_transformer_id and not individual_transformer_id:
         raise ValueError("At least one transformer id is required in job config")
 
-    households_to_process = batch.household_set.filter(removed=False).only("pk", "flex_fields")
-    individuals_to_process = batch.individual_set.filter(removed=False).only("pk", "flex_fields")
-    transformed_households = 0
-    transformed_individuals = 0
-
-    if batch.program.is_master_detail and household_transformer:
-        for household in households_to_process.iterator():
-            transformed_households += int(_apply_transformer_only(household, household_transformer))
-
-        _sync_household_refs(batch)
-
-    if individual_transformer:
-        for individual in individuals_to_process.iterator():
-            transformed_individuals += int(_apply_transformer_only(individual, individual_transformer))
-        sync_collector_links(individuals_to_process)
+    postprocessing_result = run_batch_postprocessing(
+        batch,
+        household_transformer_id=household_transformer_id,
+        individual_transformer_id=individual_transformer_id,
+        sync_household_refs=_sync_household_refs,
+    )
 
     response = {
         "batch_id": batch_id,
         "batch_name": batch.name,
-        "transformed_individuals": transformed_individuals,
+        "transformed_individuals": postprocessing_result.get("transformed_individuals", 0),
     }
     if batch.program.is_master_detail:
-        response["transformed_households"] = transformed_households
+        response["transformed_households"] = postprocessing_result.get("transformed_households", 0)
 
     logger.info("Batch transformer apply finished: %s", response)
     return response

tests/workspace/admin/batch/test_batch_reprocessing.py

@@ -12,6 +12,7 @@
     _sync_household_refs,
     _sync_kobo_household_refs,
     _sync_rdi_household_refs,
+    apply_batch_transformers,
     reprocess_batch,
 )
 from country_workspace.workspaces.models import CountryBatch
@@ -447,3 +448,78 @@ def test_reprocess_batch_creates_validation_jobs_for_people_only_batch(
     assert list(validation_jobs.call_args.kwargs["queryset"]) == list(batch.individual_set.filter(removed=False))
 
     postprocessing.assert_called_once()
+
+
+# --- apply_batch_transformers ----------------------------------------------------
+
+
+def test_apply_batch_transformers_requires_batch_id(user: User) -> None:
+    from testutils.factories import AsyncJobFactory
+
+    job = AsyncJobFactory(owner=user, config={})
+
+    with pytest.raises(ValueError, match="batch_id is required"):
+        apply_batch_transformers(job)
+
+
+def test_apply_batch_transformers_requires_transformer_ids(batch: CountryBatch, user: User) -> None:
+    from testutils.factories import AsyncJobFactory
+
+    job = AsyncJobFactory(owner=user, batch=batch, program=batch.program, config={"batch_id": batch.pk})
+
+    with pytest.raises(ValueError, match="At least one transformer id is required"):
+        apply_batch_transformers(job)
+
+
+def test_apply_batch_transformers_raises_for_missing_transformer(batch: CountryBatch, user: User) -> None:
+    from testutils.factories import AsyncJobFactory
+
+    job = AsyncJobFactory(
+        owner=user,
+        batch=batch,
+        program=batch.program,
+        config={"batch_id": batch.pk, "individual_transformer_id": 999999},
+    )
+
+    with pytest.raises(ValueError, match="Individual transformer 999999 is not available for this batch"):
+        apply_batch_transformers(job)
+
+
+def test_apply_batch_transformers_delegates_to_postprocessing(
+    batch: CountryBatch,
+    user: User,
+    mocker,
+) -> None:
+    from testutils.factories import AsyncJobFactory, TransformerFactory
+
+    household_transformer = TransformerFactory(office=batch.country_office)
+    individual_transformer = TransformerFactory(office=batch.country_office)
+    job = AsyncJobFactory(
+        owner=user,
+        batch=batch,
+        program=batch.program,
+        config={
+            "batch_id": batch.pk,
+            "household_transformer_id": household_transformer.pk,
+            "individual_transformer_id": individual_transformer.pk,
+        },
+    )
+
+    postprocessing = mocker.patch(
+        "country_workspace.workspaces.admin.batch.reprocessing.run_batch_postprocessing",
+        return_value={"transformed_households": 3, "transformed_individuals": 7, "collector_links": 5},
+    )
+
+    response = apply_batch_transformers(job)
+
+    postprocessing.assert_called_once_with(
+        batch,
+        household_transformer_id=household_transformer.pk,
+        individual_transformer_id=individual_transformer.pk,
+        sync_household_refs=_sync_household_refs,
+    )
+    assert response["batch_id"] == batch.pk
+    assert response["batch_name"] == batch.name
+    assert response["transformed_individuals"] == 7
+    if batch.program.is_master_detail:
+        assert response["transformed_households"] == 3