feat(data_classes): add API Gateway Websocket event (#6287)

* feat(data_classes): add API Gateway Websocket event

Pydantic models were added for API Gateway WebSocket events here:
https://docs.powertools.aws.dev/lambda/python/latest/utilities/parser/#built-in-models

Add the corresponding data model class.

* feat(data_classes): increase tests code coverage

---------

Co-authored-by: Leandro Damascena <[email protected]>

Author: ericbn

Diff for: aws_lambda_powertools/utilities/data_classes/__init__.py

@@ -4,6 +4,7 @@
 
 from .alb_event import ALBEvent
 from .api_gateway_proxy_event import APIGatewayProxyEvent, APIGatewayProxyEventV2
+from .api_gateway_websocket_event import APIGatewayWebSocketEvent
 from .appsync_resolver_event import AppSyncResolverEvent
 from .aws_config_rule_event import AWSConfigRuleEvent
 from .bedrock_agent_event import BedrockAgentEvent
@@ -51,6 +52,7 @@
 __all__ = [
     "APIGatewayProxyEvent",
     "APIGatewayProxyEventV2",
+    "APIGatewayWebSocketEvent",
     "SecretsManagerEvent",
     "AppSyncResolverEvent",
     "ALBEvent",

Diff for: aws_lambda_powertools/utilities/data_classes/api_gateway_websocket_event.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+import base64
+from functools import cached_property
+from typing import Any
+
+from aws_lambda_powertools.utilities.data_classes.common import (
+    CaseInsensitiveDict,
+    DictWrapper,
+)
+
+
+class APIGatewayWebSocketEventIdentity(DictWrapper):
+    @property
+    def source_ip(self) -> str:
+        return self["sourceIp"]
+
+    @property
+    def user_agent(self) -> str | None:
+        return self.get("userAgent")
+
+
+class APIGatewayWebSocketEventRequestContext(DictWrapper):
+    @property
+    def route_key(self) -> str:
+        return self["routeKey"]
+
+    @property
+    def disconnect_status_code(self) -> int | None:
+        return self.get("disconnectStatusCode")
+
+    @property
+    def message_id(self) -> str | None:
+        return self.get("messageId")
+
+    @property
+    def event_type(self) -> str:
+        return self["eventType"]
+
+    @property
+    def extended_request_id(self) -> str:
+        return self["extendedRequestId"]
+
+    @property
+    def request_time(self) -> str:
+        return self["requestTime"]
+
+    @property
+    def message_direction(self) -> str:
+        return self["messageDirection"]
+
+    @property
+    def disconnect_reason(self) -> str | None:
+        return self.get("disconnectReason")
+
+    @property
+    def stage(self) -> str:
+        return self["stage"]
+
+    @property
+    def connected_at(self) -> int:
+        return self["connectedAt"]
+
+    @property
+    def request_time_epoch(self) -> int:
+        return self["requestTimeEpoch"]
+
+    @property
+    def identity(self) -> APIGatewayWebSocketEventIdentity:
+        return APIGatewayWebSocketEventIdentity(self["identity"])
+
+    @property
+    def request_id(self) -> str:
+        return self["requestId"]
+
+    @property
+    def domain_name(self) -> str:
+        return self["domainName"]
+
+    @property
+    def connection_id(self) -> str:
+        return self["connectionId"]
+
+    @property
+    def api_id(self) -> str:
+        return self["apiId"]
+
+
+class APIGatewayWebSocketEvent(DictWrapper):
+    """AWS proxy integration event for WebSocket API
+
+    Documentation:
+    --------------
+    - https://docs.aws.amazon.com/apigateway/latest/developerguide/apigateway-websocket-api-integration-requests.html
+    """
+
+    @property
+    def is_base64_encoded(self) -> bool:
+        return self["isBase64Encoded"]
+
+    @property
+    def body(self) -> str | None:
+        return self.get("body")
+
+    @cached_property
+    def decoded_body(self) -> str | None:
+        body = self.body
+        if self.is_base64_encoded and body:
+            return base64.b64decode(body.encode()).decode()
+        return body
+
+    @cached_property
+    def json_body(self) -> Any:
+        if self.decoded_body:
+            return self._json_deserializer(self.decoded_body)
+        return None
+
+    @property
+    def headers(self) -> dict[str, str]:
+        return CaseInsensitiveDict(self.get("headers"))
+
+    @property
+    def multi_value_headers(self) -> dict[str, list[str]]:
+        return CaseInsensitiveDict(self.get("multiValueHeaders"))
+
+    @property
+    def request_context(self) -> APIGatewayWebSocketEventRequestContext:
+        return APIGatewayWebSocketEventRequestContext(self["requestContext"])

Diff for: tests/unit/data_classes/required_dependencies/test_api_gateway_websocket_event.py

@@ -0,0 +1,109 @@
+import json
+
+from aws_lambda_powertools.utilities.data_classes import APIGatewayWebSocketEvent
+from tests.functional.utils import load_event
+
+
+def test_connect_api_gateway_websocket_event():
+    raw_event = load_event("apiGatewayWebSocketApiConnect.json")
+    parsed_event = APIGatewayWebSocketEvent(raw_event)
+
+    assert parsed_event.is_base64_encoded is False
+    assert parsed_event.body is None
+    assert parsed_event.decoded_body is None
+    assert parsed_event.json_body is None
+    assert parsed_event.headers == raw_event["headers"]
+    assert parsed_event.multi_value_headers == raw_event["multiValueHeaders"]
+
+    request_context = parsed_event.request_context
+    request_context_raw = raw_event["requestContext"]
+    assert request_context.route_key == request_context_raw["routeKey"]
+    assert request_context.disconnect_status_code is None
+    assert request_context.message_id is None
+    assert request_context.event_type == request_context_raw["eventType"]
+    assert request_context.extended_request_id == request_context_raw["extendedRequestId"]
+    assert request_context.request_time == request_context_raw["requestTime"]
+    assert request_context.message_direction == request_context_raw["messageDirection"]
+    assert request_context.disconnect_reason is None
+    assert request_context.stage == request_context_raw["stage"]
+    assert request_context.connected_at == request_context_raw["connectedAt"]
+    assert request_context.request_time_epoch == request_context_raw["requestTimeEpoch"]
+    assert request_context.request_id == request_context_raw["requestId"]
+    assert request_context.domain_name == request_context_raw["domainName"]
+    assert request_context.connection_id == request_context_raw["connectionId"]
+    assert request_context.api_id == request_context_raw["apiId"]
+
+    identity = request_context.identity
+    identity_raw = request_context_raw["identity"]
+    assert identity.source_ip == identity_raw["sourceIp"]
+    assert identity.user_agent is None
+
+
+def test_disconnect_api_gateway_websocket_event():
+    raw_event = load_event("apiGatewayWebSocketApiDisconnect.json")
+    parsed_event = APIGatewayWebSocketEvent(raw_event)
+
+    assert parsed_event.is_base64_encoded is False
+    assert parsed_event.body is None
+    assert parsed_event.decoded_body is None
+    assert parsed_event.json_body is None
+    assert parsed_event.headers == raw_event["headers"]
+    assert parsed_event.multi_value_headers == raw_event["multiValueHeaders"]
+
+    request_context = parsed_event.request_context
+    request_context_raw = raw_event["requestContext"]
+    assert request_context.route_key == request_context_raw["routeKey"]
+    assert request_context.disconnect_status_code == request_context_raw["disconnectStatusCode"]
+    assert request_context.message_id is None
+    assert request_context.event_type == request_context_raw["eventType"]
+    assert request_context.extended_request_id == request_context_raw["extendedRequestId"]
+    assert request_context.request_time == request_context_raw["requestTime"]
+    assert request_context.message_direction == request_context_raw["messageDirection"]
+    assert request_context.disconnect_reason == request_context_raw["disconnectReason"]
+    assert request_context.stage == request_context_raw["stage"]
+    assert request_context.connected_at == request_context_raw["connectedAt"]
+    assert request_context.request_time_epoch == request_context_raw["requestTimeEpoch"]
+    assert request_context.request_id == request_context_raw["requestId"]
+    assert request_context.domain_name == request_context_raw["domainName"]
+    assert request_context.connection_id == request_context_raw["connectionId"]
+    assert request_context.api_id == request_context_raw["apiId"]
+
+    identity = request_context.identity
+    identity_raw = request_context_raw["identity"]
+    assert identity.source_ip == identity_raw["sourceIp"]
+    assert identity.user_agent is None
+
+
+def test_message_api_gateway_websocket_event():
+    raw_event = load_event("apiGatewayWebSocketApiMessage.json")
+    parsed_event = APIGatewayWebSocketEvent(raw_event)
+
+    assert parsed_event.is_base64_encoded is False
+    assert parsed_event.body == raw_event["body"]
+    assert parsed_event.decoded_body == raw_event["body"]
+    assert parsed_event.json_body == json.loads(raw_event["body"])
+    assert parsed_event.headers == {}
+    assert parsed_event.multi_value_headers == {}
+
+    request_context = parsed_event.request_context
+    request_context_raw = raw_event["requestContext"]
+    assert request_context.route_key == request_context_raw["routeKey"]
+    assert request_context.disconnect_status_code is None
+    assert request_context.message_id == request_context_raw["messageId"]
+    assert request_context.event_type == request_context_raw["eventType"]
+    assert request_context.extended_request_id == request_context_raw["extendedRequestId"]
+    assert request_context.request_time == request_context_raw["requestTime"]
+    assert request_context.message_direction == request_context_raw["messageDirection"]
+    assert request_context.disconnect_reason is None
+    assert request_context.stage == request_context_raw["stage"]
+    assert request_context.connected_at == request_context_raw["connectedAt"]
+    assert request_context.request_time_epoch == request_context_raw["requestTimeEpoch"]
+    assert request_context.request_id == request_context_raw["requestId"]
+    assert request_context.domain_name == request_context_raw["domainName"]
+    assert request_context.connection_id == request_context_raw["connectionId"]
+    assert request_context.api_id == request_context_raw["apiId"]
+
+    identity = request_context.identity
+    identity_raw = request_context_raw["identity"]
+    assert identity.source_ip == identity_raw["sourceIp"]
+    assert identity.user_agent is None