From c5d996329d5b0368ecbda7ebabca97f38b3d168d Mon Sep 17 00:00:00 2001 From: Jaime Hablutzel Date: Wed, 9 Apr 2025 15:58:11 -0500 Subject: [PATCH] Improve error handling for perspectives Lambda errors. --- .../mpic_coordinator_lambda_function.py | 12 ++-- .../test_mpic_coordinator_lambda.py | 59 ++++++++++++++----- 2 files changed, 51 insertions(+), 20 deletions(-) diff --git a/src/aws_lambda_mpic/mpic_coordinator_lambda/mpic_coordinator_lambda_function.py b/src/aws_lambda_mpic/mpic_coordinator_lambda/mpic_coordinator_lambda_function.py index 8d3ade9..50e2e82 100644 --- a/src/aws_lambda_mpic/mpic_coordinator_lambda/mpic_coordinator_lambda_function.py +++ b/src/aws_lambda_mpic/mpic_coordinator_lambda/mpic_coordinator_lambda_function.py @@ -31,6 +31,10 @@ class PerspectiveEndpoints(BaseModel): caa_endpoint_info: PerspectiveEndpointInfo +class LambdaExecutionException(Exception): + pass + + class MpicCoordinatorLambdaHandler: def __init__(self): perspectives_json = os.environ["perspectives"] @@ -140,11 +144,11 @@ async def call_remote_perspective( InvocationType="RequestResponse", Payload=check_request.model_dump_json(), # AWS Lambda functions expect a JSON string for payload ) - response_payload = json.loads(await response["Payload"].read()) + response_payload = await response["Payload"].read() + if 'FunctionError' in response: + raise LambdaExecutionException(f"Lambda execution error: {response_payload.decode('utf-8')}") + response_payload = json.loads(response_payload) return self.check_response_adapter.validate_json(response_payload["body"]) - except ValidationError as ve: - self.logger.error(msg=f"Validation error in response from {perspective.code}: {ve}") - raise ve finally: await self.release_lambda_client(perspective.code, client) diff --git a/tests/unit/aws_lambda_mpic/test_mpic_coordinator_lambda.py b/tests/unit/aws_lambda_mpic/test_mpic_coordinator_lambda.py index d3a4800..fa098e7 100644 --- a/tests/unit/aws_lambda_mpic/test_mpic_coordinator_lambda.py +++ b/tests/unit/aws_lambda_mpic/test_mpic_coordinator_lambda.py @@ -11,7 +11,7 @@ APIGatewayEventRequestContext, APIGatewayEventIdentity, ) -from pydantic import TypeAdapter +from pydantic import TypeAdapter, BaseModel from open_mpic_core import RemotePerspective from open_mpic_core import DcvCheckRequest, DcvCheckResponse, CaaCheckResponse, PerspectiveResponse @@ -22,8 +22,9 @@ from open_mpic_core import MpicCaaResponse from botocore.response import StreamingBody import aws_lambda_mpic.mpic_coordinator_lambda.mpic_coordinator_lambda_function as mpic_coordinator_lambda_function -from aws_lambda_mpic.mpic_coordinator_lambda.mpic_coordinator_lambda_function import MpicCoordinatorLambdaHandler from aws_lambda_mpic.mpic_coordinator_lambda.mpic_coordinator_lambda_function import ( + MpicCoordinatorLambdaHandler, + LambdaExecutionException, PerspectiveEndpoints, PerspectiveEndpointInfo, ) @@ -78,22 +79,9 @@ def set_env_variables(): async def call_remote_perspective__should_make_aws_lambda_call_with_provided_arguments_and_return_check_response( self, set_env_variables, mocker ): - # Mock the aioboto3 client creation and context manager - mock_client = AsyncMock() - mock_client.invoke = AsyncMock(side_effect=self.create_successful_aioboto3_response_for_dcv_check) - - # Mock the __aenter__ method that gets called in initialize_client_pools() - mock_client.__aenter__.return_value = mock_client - - # Mock the session creation and client initialization - mock_session = mocker.patch("aioboto3.Session") - mock_session.return_value.client.return_value = mock_client + lambda_handler, mock_client = await self.mock_lambda_handler_for_lambda_invoke(mocker, self.create_successful_aioboto3_response_for_dcv_check) dcv_check_request = ValidCheckCreator.create_valid_dns_check_request() - lambda_handler = MpicCoordinatorLambdaHandler() - - await lambda_handler.initialize_client_pools() - perspective_code = "us-west-1" check_response = await lambda_handler.call_remote_perspective( RemotePerspective(code=perspective_code, rir="arin"), CheckType.DCV, dcv_check_request @@ -111,6 +99,20 @@ async def call_remote_perspective__should_make_aws_lambda_call_with_provided_arg Payload=dcv_check_request.model_dump_json(), ) + async def call_remote_perspective__should_make_aws_lambda_call_and_handle_lambda_execution_exceptions( + self, set_env_variables, mocker + ): + lambda_handler, mock_client = await self.mock_lambda_handler_for_lambda_invoke(mocker, self.create_error_aioboto3_response) + + class Dummy(BaseModel): + pass + + with pytest.raises(LambdaExecutionException) as exc_info: + await lambda_handler.call_remote_perspective( + RemotePerspective(code="us-west-1", rir="dummy"), CheckType.DCV, Dummy() + ) + assert exc_info.value.args[0] == "Lambda execution error: {\"errorMessage\": \"some message\"}" + def lambda_handler__should_return_400_error_and_details_given_invalid_request_body(self): request = ValidMpicRequestCreator.create_valid_dcv_mpic_request() # noinspection PyTypeChecker @@ -250,6 +252,19 @@ async def read(self): return json_bytes return {"Payload": MockStreamingBody()} + + # noinspection PyUnusedLocal + async def create_error_aioboto3_response(self, *args, **kwargs): + expected_response = {"errorMessage": "some message"} + json_bytes = json.dumps(expected_response).encode("utf-8") + + class MockStreamingBody: + # noinspection PyMethodMayBeStatic + async def read(self): + return json_bytes + + # See https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/lambda/client/invoke.html, "Response Structure" + return {"Payload": MockStreamingBody(), "FunctionError": None} @staticmethod def create_caa_mpic_response(): @@ -304,6 +319,18 @@ def get_perspectives_by_code_dict_from_file() -> dict[str, RemotePerspective]: perspectives = perspective_type_adapter.validate_python(perspectives_yaml["aws_available_regions"]) return {perspective.code: perspective for perspective in perspectives} + async def mock_lambda_handler_for_lambda_invoke(self, mocker, lambda_invoke_side_effect): + # Mock the aioboto3 client creation and context manager + mock_client = AsyncMock() + mock_client.invoke = AsyncMock(side_effect=lambda_invoke_side_effect) + # Mock the __aenter__ method that gets called in initialize_client_pools() + mock_client.__aenter__.return_value = mock_client + # Mock the session creation and client initialization + mock_session = mocker.patch("aioboto3.Session") + mock_session.return_value.client.return_value = mock_client + lambda_handler = MpicCoordinatorLambdaHandler() + await lambda_handler.initialize_client_pools() + return lambda_handler, mock_client if __name__ == "__main__": pytest.main()