Adding more tests

Author: leandrodamascena

aws_lambda_powertools/event_handler/openapi/models.py

@@ -4,6 +4,7 @@
 from pydantic import AnyUrl, BaseModel, Field
 
 from aws_lambda_powertools.event_handler.openapi.compat import model_rebuild, parser_openapi_extension
+from aws_lambda_powertools.event_handler.openapi.exceptions import SchemaValidationError
 from aws_lambda_powertools.event_handler.openapi.pydantic_loader import PYDANTIC_V2
 from aws_lambda_powertools.shared.types import Annotated, Literal
 
@@ -25,6 +26,7 @@ class OpenAPIExtensions(BaseModel):
 
     # This rule is valid for Pydantic v1 and v2
     # If the 'openapi_extensions' field is present in the 'values' dictionary,
+    # And if the extension starts with x-
     # update the 'values' dictionary with the contents of 'openapi_extensions',
     # and then remove the 'openapi_extensions' field from the 'values' dictionary
 
@@ -34,8 +36,15 @@ class OpenAPIExtensions(BaseModel):
 
         @parser_openapi_extension(mode="before")
         def serialize_openapi_extension_v2(self):
-            if isinstance(self, dict) and self.get("openapi_extensions"):
-                self.update(self.get("openapi_extensions"))
+            openapi_extension_value = self.get("openapi_extensions")
+
+            if isinstance(self, dict) and openapi_extension_value:
+
+                for extension_key in openapi_extension_value:
+                    if not str(extension_key).startswith("x-"):
+                        raise SchemaValidationError("An OpenAPI extension key must start with x-")
+
+                self.update(openapi_extension_value)
                 self.pop("openapi_extensions", None)
 
             return self
@@ -44,9 +53,17 @@ def serialize_openapi_extension_v2(self):
 
         @parser_openapi_extension(pre=False, allow_reuse=True)
         def serialize_openapi_extension_v1(cls, values):
-            if values.get("openapi_extensions"):
+            openapi_extension_value = values.get("openapi_extensions")
+
+            if openapi_extension_value:
+
+                for extension_key in openapi_extension_value:
+                    if not str(extension_key).startswith("x-"):
+                        raise SchemaValidationError("An OpenAPI extension key must start with x-")
+
                 values.update(values["openapi_extensions"])
                 del values["openapi_extensions"]
+
             return values
 
         class Config:

tests/unit/event_handler/_pydantic/test_openapi_models_pydantic_v1.py

@@ -1,5 +1,6 @@
 import pytest
 
+from aws_lambda_powertools.event_handler.openapi.exceptions import SchemaValidationError
 from aws_lambda_powertools.event_handler.openapi.models import OpenAPIExtensions
 
 
@@ -12,6 +13,14 @@ def test_openapi_extensions_with_dict():
     assert extensions.dict(exclude_none=True) == {"x-amazon-apigateway": {"foo": "bar"}}
 
 
+@pytest.mark.usefixtures("pydanticv1_only")
+def test_openapi_extensions_with_invalid_key():
+    # GIVEN we create an OpenAPIExtensions object with an invalid value
+    with pytest.raises(SchemaValidationError):
+        # THEN must raise an exception
+        OpenAPIExtensions(openapi_extensions={"amazon-apigateway-invalid": {"foo": "bar"}})
+
+
 @pytest.mark.usefixtures("pydanticv1_only")
 def test_openapi_extensions_with_proxy_models():
 

tests/unit/event_handler/_pydantic/test_openapi_models_pydantic_v2.py

@@ -1,5 +1,6 @@
 import pytest
 
+from aws_lambda_powertools.event_handler.openapi.exceptions import SchemaValidationError
 from aws_lambda_powertools.event_handler.openapi.models import OpenAPIExtensions
 
 
@@ -12,6 +13,14 @@ def test_openapi_extensions_with_dict():
     assert extensions.model_dump(exclude_none=True) == {"x-amazon-apigateway": {"foo": "bar"}}
 
 
+@pytest.mark.usefixtures("pydanticv2_only")
+def test_openapi_extensions_with_invalid_key():
+    # GIVEN we create an OpenAPIExtensions object with an invalid value
+    with pytest.raises(SchemaValidationError):
+        # THEN must raise an exception
+        OpenAPIExtensions(openapi_extensions={"amazon-apigateway-invalid": {"foo": "bar"}})
+
+
 @pytest.mark.usefixtures("pydanticv2_only")
 def test_openapi_extensions_with_proxy_models():