diff --git a/src/aws_durable_execution_sdk_python_testing/web/handlers.py b/src/aws_durable_execution_sdk_python_testing/web/handlers.py index 7c42d09..85417e8 100644 --- a/src/aws_durable_execution_sdk_python_testing/web/handlers.py +++ b/src/aws_durable_execution_sdk_python_testing/web/handlers.py @@ -2,6 +2,7 @@ from __future__ import annotations +import base64 import json import logging from abc import ABC, abstractmethod @@ -183,6 +184,37 @@ def _no_content_response( # Removed deprecated _error_response method - use AWS exceptions directly + def _parse_callback_result_payload(self, request: HTTPRequest) -> bytes: + """Parse callback result payload from request body. + + Expects JSON payload with base64-encoded Result field. + + Args: + request: The HTTP request containing the JSON payload + + Returns: + bytes: The decoded result payload + + Raises: + InvalidParameterValueException: If payload parsing fails + """ + if not request.body or not isinstance(request.body, bytes): + return b"" + + try: + payload = json.loads(request.body.decode("utf-8")) + if isinstance(payload, dict) and "Result" in payload: + result_value = payload["Result"] + if isinstance(result_value, str): + return base64.b64decode(result_value) + return b"" + except (json.JSONDecodeError, UnicodeDecodeError) as e: + msg = f"Failed to parse JSON payload: {e}" + raise InvalidParameterValueException(msg) from e + except ValueError as e: + msg = f"Failed to decode base64 result: {e}" + raise InvalidParameterValueException(msg) from e + def _parse_query_param(self, request: HTTPRequest, param_name: str) -> str | None: """Parse a single query parameter from the request. @@ -611,8 +643,7 @@ def handle(self, parsed_route: Route, request: HTTPRequest) -> HTTPResponse: callback_route = cast(CallbackSuccessRoute, parsed_route) callback_id: str = callback_route.callback_id - # For binary payload operations, body is raw bytes - result_bytes = request.body if isinstance(request.body, bytes) else b"" + result_bytes: bytes = self._parse_callback_result_payload(request) callback_response: SendDurableExecutionCallbackSuccessResponse = ( # noqa: F841 self.executor.send_callback_success( diff --git a/tests/web/handlers_test.py b/tests/web/handlers_test.py index 5cc4fb0..52c2dea 100644 --- a/tests/web/handlers_test.py +++ b/tests/web/handlers_test.py @@ -2,6 +2,8 @@ from __future__ import annotations +import base64 +import json from typing import TYPE_CHECKING, Any from unittest.mock import Mock @@ -2037,13 +2039,15 @@ def test_send_durable_execution_callback_success_handler(): assert isinstance(route, CallbackSuccessRoute) assert route.callback_id == "test-callback-id" - # Test with valid request body (bytes for callback operations) + result_data = base64.b64encode(b"success-result").decode("utf-8") + request_body = json.dumps({"Result": result_data}).encode("utf-8") + request = HTTPRequest( method="POST", path=route, headers={"Content-Type": "application/json"}, query_params={}, - body=b"success-result", + body=request_body, ) response = handler.handle(route, request) @@ -2058,6 +2062,57 @@ def test_send_durable_execution_callback_success_handler(): ) +def test_send_durable_execution_callback_success_handler_invalid_json(): + """Test SendDurableExecutionCallbackSuccessHandler with invalid JSON.""" + executor = Mock() + handler = SendDurableExecutionCallbackSuccessHandler(executor) + + router = Router() + route = router.find_route( + "/2025-12-01/durable-execution-callbacks/test-callback-id/succeed", "POST" + ) + + request = HTTPRequest( + method="POST", + path=route, + headers={"Content-Type": "application/json"}, + query_params={}, + body=b"invalid-json", + ) + + response = handler.handle(route, request) + + assert response.status_code == 400 + assert response.body["Type"] == "InvalidParameterValueException" + assert "Failed to parse JSON payload" in response.body["message"] + + +def test_send_durable_execution_callback_success_handler_invalid_base64(): + """Test SendDurableExecutionCallbackSuccessHandler with invalid base64.""" + executor = Mock() + handler = SendDurableExecutionCallbackSuccessHandler(executor) + + router = Router() + route = router.find_route( + "/2025-12-01/durable-execution-callbacks/test-callback-id/succeed", "POST" + ) + + request_body = json.dumps({"Result": "invalid-base64!"}).encode("utf-8") + request = HTTPRequest( + method="POST", + path=route, + headers={"Content-Type": "application/json"}, + query_params={}, + body=request_body, + ) + + response = handler.handle(route, request) + + assert response.status_code == 400 + assert response.body["Type"] == "InvalidParameterValueException" + assert "Failed to decode base64 result" in response.body["message"] + + def test_send_durable_execution_callback_success_handler_empty_body(): """Test SendDurableExecutionCallbackSuccessHandler with empty body.""" executor = Mock()