diff --git a/packages/aws-durable-execution-sdk-python/tests/e2e/checkpoint_response_int_test.py b/packages/aws-durable-execution-sdk-python/tests/e2e/checkpoint_response_int_test.py index 1fb5d21a..5ecd0123 100644 --- a/packages/aws-durable-execution-sdk-python/tests/e2e/checkpoint_response_int_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/e2e/checkpoint_response_int_test.py @@ -11,9 +11,13 @@ import pytest +from aws_durable_execution_sdk_python.concurrency.models import BatchResult from aws_durable_execution_sdk_python.config import ChildConfig, Duration from aws_durable_execution_sdk_python.context import DurableContext, durable_step -from aws_durable_execution_sdk_python.exceptions import InvocationError +from aws_durable_execution_sdk_python.exceptions import ( + ChildContextError, + InvocationError, +) from aws_durable_execution_sdk_python.execution import ( InvocationStatus, durable_execution, @@ -762,3 +766,75 @@ def my_handler(event, context: DurableContext) -> str: if hasattr(op, "action") and op.action.value == "FAIL" ] assert len(fail_updates) == 0 + + +def test_parallel_and_child_context_agree_on_custom_error_type(): + """ctx.parallel() and ctx.run_in_child_context() surface the same + error_type for an identical custom exception: the caller's original type, + not the ChildContextError wrapper class name. + """ + + class PermanentFailure(Exception): + pass + + def branch(child_context: DurableContext) -> str: + msg: str = "Invalid input data" + raise PermanentFailure(msg) + + captured: dict[str, str | None] = {} + + @durable_execution + def my_handler(event, context: DurableContext) -> str: + try: + context.run_in_child_context(branch) + except ChildContextError as e: + captured["child"] = e.error_type + + result: BatchResult[str] = context.parallel([branch]) + try: + result.throw_if_error() + except ChildContextError as e: + captured["parallel"] = e.error_type + + return "handled" + + with patch( + "aws_durable_execution_sdk_python.execution.LambdaClient" + ) as mock_client_class: + mock_client = Mock() + mock_client_class.initialize_client.return_value = mock_client + + mock_checkpoint, _ = create_mock_checkpoint_with_operations() + mock_client.checkpoint = mock_checkpoint + + event = { + "DurableExecutionArn": "test-arn/execution-1", + "CheckpointToken": "test-token", + "InitialExecutionState": { + "Operations": [ + { + "Id": "execution-1", + "Type": "EXECUTION", + "Status": "STARTED", + "ExecutionDetails": {"InputPayload": "{}"}, + } + ], + "NextMarker": "", + }, + "LocalRunner": True, + } + + lambda_context = Mock() + lambda_context.aws_request_id = "test-request-id" + lambda_context.client_context = None + lambda_context.identity = None + lambda_context._epoch_deadline_time_in_ms = 0 # noqa: SLF001 + lambda_context.invoked_function_arn = "test-arn" + lambda_context.tenant_id = None + + my_handler(event, lambda_context) + + # Both operations surface the caller's original type, and they agree. + assert captured["child"] == "PermanentFailure" + assert captured["parallel"] == "PermanentFailure" + assert captured["child"] == captured["parallel"] diff --git a/packages/aws-durable-execution-sdk-python/tests/operation/child_test.py b/packages/aws-durable-execution-sdk-python/tests/operation/child_test.py index ec06dfe0..07c7150e 100644 --- a/packages/aws-durable-execution-sdk-python/tests/operation/child_test.py +++ b/packages/aws-durable-execution-sdk-python/tests/operation/child_test.py @@ -15,6 +15,7 @@ DurableApiErrorCategory, ExecutionError, InvocationError, + StepError, ) from aws_durable_execution_sdk_python.identifier import OperationIdentifier from aws_durable_execution_sdk_python.lambda_service import ( @@ -472,6 +473,75 @@ def test_child_handler_non_retryable_invocation_error_wrapped(): assert fail_operation.action is OperationAction.FAIL +def test_child_handler_preserves_data_across_nested_boundaries(): + """Error data and stack_trace survive >=2 nested run_in_child_context + boundaries. + + Each boundary carries the escaping error's data and stack_trace forward, + both on the raised ChildContextError and on the FAIL checkpoint that a + replay reconstructs from. + """ + mock_state = Mock(spec=ExecutionState) + mock_state.durable_execution_arn = "test_arn" + mock_result = Mock() + mock_result.is_succeeded.return_value = False + mock_result.is_failed.return_value = False + mock_result.is_started.return_value = False + mock_result.is_replay_children.return_value = False + mock_result.is_existent.return_value = False + mock_state.get_checkpoint_result.return_value = mock_result + # Return each function unchanged so both nested bodies actually run. + mock_state.wrap_user_function.side_effect = lambda func, *args, **kwargs: func + + # Innermost failure carries a serialized payload and stack trace, as an + # operation error reconstructed from a checkpoint would. + inner_error: StepError = StepError( + "inner step failed", + error_type="ValueError", + data="serialized-payload", + stack_trace=["frame-a", "frame-b"], + ) + + def inner_body(): + raise inner_error + + def outer_body(): + return child_handler( + inner_body, + mock_state, + OperationIdentifier( + "inner", OperationSubType.RUN_IN_CHILD_CONTEXT, None, "inner_name" + ), + None, + ) + + with pytest.raises(ChildContextError) as exc_info: + child_handler( + outer_body, + mock_state, + OperationIdentifier( + "outer", OperationSubType.RUN_IN_CHILD_CONTEXT, None, "outer_name" + ), + None, + ) + + # Data and stack_trace survive both boundaries on the surfaced error. + assert exc_info.value.data == "serialized-payload" + assert exc_info.value.stack_trace == ["frame-a", "frame-b"] + + # Both FAIL checkpoints (inner and outer) recorded the same payload, so the + # value also survives a replay that rebuilds from either checkpoint. + fail_operations = [ + call.kwargs["operation_update"] + for call in mock_state.create_checkpoint.call_args_list + if call.kwargs["operation_update"].action is OperationAction.FAIL + ] + assert len(fail_operations) == 2 + for fail_operation in fail_operations: + assert fail_operation.error.data == "serialized-payload" + assert fail_operation.error.stack_trace == ["frame-a", "frame-b"] + + def test_child_handler_with_config(): """Test child_handler with config parameter.""" mock_state = Mock(spec=ExecutionState)