diff --git a/tests/utilities/aws_lambda/test_task_handler.py b/tests/utilities/aws_lambda/test_task_handler.py index fae35de28..418669bcd 100644 --- a/tests/utilities/aws_lambda/test_task_handler.py +++ b/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: diff --git a/utilities/aws_lambda/task_handler.py b/utilities/aws_lambda/task_handler.py index 43699aee4..0ed680c92 100644 --- a/utilities/aws_lambda/task_handler.py +++ b/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.