task_handler passes orchestrator and task_id to wrapped function when flag is true 🟥

tests/utilities/aws_lambda/test_task_handler.py

@@ -77,6 +77,33 @@ def test_task_handler_records_cloudwatch_url_on_subtask(
     assert "$255B$2524LATEST$255D" in saved_url
 
 
+def test_task_handler_passes_orchestrator_and_task_id_when_flag_is_true(
+    harness: Harness,
+) -> None:
+    # Arrange
+    received: list[tuple[TaskOrchestrator, 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((orchestrator, task_id))
+
+    # Act
+    result = _handler(_direct_event("prop-1"), context=None)
+
+    # Assert
+    assert len(received) == 1
+    recv_orchestrator, recv_task_id = received[0]
+    assert recv_orchestrator is harness.orchestrator
+    assert recv_task_id == UUID(result[0]["task_id"])
+
+
 def test_task_handler_leaves_cloudwatch_url_unset_outside_lambda(
     harness: Harness, monkeypatch: pytest.MonkeyPatch
 ) -> None:

utilities/aws_lambda/task_handler.py

@@ -25,6 +25,7 @@ def task_handler(
     task_source: str,
     source: Source,
     orchestrator_cm: Optional[OrchestratorCM] = None,
+    pass_task_orchestrator: bool = False,
 ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
     """Run the wrapped function as the body of a freshly-created Task + SubTask.