From ca6ca2a220b32d7fbc71f5aa479db277882f697b Mon Sep 17 00:00:00 2001
From: Khalim Conn-Kowlessar <kconnkowlessar@gmail.com>
Date: Tue, 7 Jul 2026 11:47:20 +0000
Subject: [PATCH] =?UTF-8?q?A=20batch=20message=20with=20task=5Fid=20and=20?=
 =?UTF-8?q?subtask=5Fid=20attaches=20to=20the=20app-owned=20task=20?=
 =?UTF-8?q?=F0=9F=9F=A5?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 .../utilities/aws_lambda/test_task_handler.py | 54 +++++++++++++++++++
 1 file changed, 54 insertions(+)

diff --git a/tests/utilities/aws_lambda/test_task_handler.py b/tests/utilities/aws_lambda/test_task_handler.py
index 418669bcd..974a47966 100644
--- a/tests/utilities/aws_lambda/test_task_handler.py
+++ b/tests/utilities/aws_lambda/test_task_handler.py
@@ -104,6 +104,60 @@ def test_task_handler_passes_orchestrator_and_task_id_when_flag_is_true(
     assert recv_task_id == UUID(result[0]["task_id"])
 
 
+def test_task_handler_attaches_to_a_supplied_task_and_subtask(
+    harness: Harness, db_engine: Engine
+) -> None:
+    """A message carrying task_id + subtask_id (a Modelling Run batch,
+    ADR-0055) runs under the distributor's pre-created sub_task — no new Task
+    or SubTask rows are created."""
+    # arrange — the distributor's pre-created task + batch sub_task
+    task, subtask = harness.orchestrator.create_task_with_subtask(
+        task_source="app:modelling_run", inputs={"property_ids": [1, 2]}
+    )
+    received: list[UUID] = []
+
+    @task_handler(
+        task_source="modelling_e2e",
+        source=Source.PROPERTY,
+        orchestrator_cm=harness.factory,
+        pass_task_orchestrator=True,
+    )
+    def handler(
+        body: dict[str, Any], context: Any, orchestrator: TaskOrchestrator, task_id: UUID
+    ) -> None:
+        received.append(task_id)
+
+    event = {
+        "Records": [
+            {
+                "messageId": "m-1",
+                "body": (
+                    f'{{"task_id": "{task.id}", "subtask_id": "{subtask.id}", '
+                    f'"property_ids": [1, 2]}}'
+                ),
+            }
+        ]
+    }
+
+    # act
+    result = handler(event, context=None)
+
+    # assert — ran under the supplied ids, sub_task completed, nothing created
+    assert received == [task.id]
+    assert result["tasks"] == [
+        {"task_id": str(task.id), "subtask_id": str(subtask.id)}
+    ]
+    assert harness.subtasks.get(subtask.id).status.value == "complete"
+    with Session(db_engine) as session:
+        from sqlmodel import select
+
+        from backend.app.db.models.tasks import SubTask as SubTaskRow
+        from backend.app.db.models.tasks import Task as TaskRow
+
+        assert len(session.exec(select(TaskRow)).all()) == 1
+        assert len(session.exec(select(SubTaskRow)).all()) == 1
+
+
 def test_task_handler_leaves_cloudwatch_url_unset_outside_lambda(
     harness: Harness, monkeypatch: pytest.MonkeyPatch
 ) -> None: