From 57e5eca3f809519d9cd1a1f68ad5aa54300fa99f Mon Sep 17 00:00:00 2001
From: rikuke <33894149+rikuke@users.noreply.github.com>
Date: Fri, 4 Oct 2024 10:29:06 +0300
Subject: [PATCH] fix: do not remove batch if ahjo automation (#3381)

---
 .../api/v1/serializers/application.py         |  6 ++++-
 .../tests/test_applications_api.py            | 22 ++++++++++++++++---
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/backend/benefit/applications/api/v1/serializers/application.py b/backend/benefit/applications/api/v1/serializers/application.py
index b5adbf01d0..353d6fd80c 100755
--- a/backend/benefit/applications/api/v1/serializers/application.py
+++ b/backend/benefit/applications/api/v1/serializers/application.py
@@ -1836,7 +1836,11 @@ def handle_status_transition(
         if instance.status == ApplicationStatus.CANCELLED:
             self._cancel_application(instance)
         self._assign_handler_if_needed(instance)
-        self._remove_batch_if_needed(instance)
+
+        if not instance.handled_by_ahjo_automation and instance.ahjo_case_id is None:
+            # If the application has been handled by the Ahjo automation, we don't want to
+            # remove the batch, as it's needed for the Ahjo automation
+            self._remove_batch_if_needed(instance)
 
     def _cancel_application(self, instance):
         instance.archived = True
diff --git a/backend/benefit/applications/tests/test_applications_api.py b/backend/benefit/applications/tests/test_applications_api.py
index 6b9fa365c5..a8ed5c5de7 100755
--- a/backend/benefit/applications/tests/test_applications_api.py
+++ b/backend/benefit/applications/tests/test_applications_api.py
@@ -1463,16 +1463,29 @@ def test_application_accept(
     assert handling_application.calculation.granted_as_de_minimis_aid is True
 
 
+@pytest.mark.parametrize(
+    "ahjo_case_id, handled_by_ahjo_automation, has_batch,",
+    [
+        (None, False, False),
+        ("KISSA123", True, True),
+    ],
+)
 def test_application_with_batch_back_to_handling(
     request,
     handler_api_client,
     decided_application,
+    ahjo_case_id,
+    handled_by_ahjo_automation,
+    has_batch,
 ):
     """
     When application is moved back to handling, the application
-    needs to be remvoved from any batch.
+    needs to be removed from any batch.
     """
-    decided_application.batch = ApplicationBatchFactory()
+    batch = ApplicationBatchFactory()
+    decided_application.batch = batch
+    decided_application.ahjo_case_id = ahjo_case_id
+    decided_application.handled_by_ahjo_automation = handled_by_ahjo_automation
     decided_application.save()
     data = HandlerApplicationSerializer(decided_application).data
     data["status"] = ApplicationStatus.HANDLING
@@ -1483,7 +1496,10 @@ def test_application_with_batch_back_to_handling(
     )
     assert response.status_code == 200
     decided_application.refresh_from_db()
-    assert decided_application.batch is None
+    if has_batch:
+        assert decided_application.batch == batch
+    else:
+        assert decided_application.batch is None
 
 
 @pytest.mark.parametrize(