Updated imports

cloutierMat · bentsku · commit a324a08dcbad · 2024-07-01T17:53:42.000+02:00
diff --git a/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/api_key_validation.py b/localstack-core/localstack/services/apigateway/next_gen/execute_api/handlers/api_key_validation.py
@@ -5,9 +5,10 @@
 from localstack.http import Response
 
 from ..api import RestApiGatewayHandler, RestApiGatewayHandlerChain
-from ..context import IdentityContext, InvocationRequest, RestApiInvocationContext
+from ..context import InvocationRequest, RestApiInvocationContext
 from ..gateway_response import InvalidAPIKeyError
 from ..moto_helpers import get_api_key, get_usage_plan_keys, get_usage_plans
+from ..variables import ContextVarsIdentity
 
 LOG = logging.getLogger(__name__)
 
@@ -36,7 +37,7 @@ def __call__(
             return
 
         # If the Identity context was not created yet, instantiate it and attach it to the context variables
-        if not (identity := context.context_variables.get("identity", IdentityContext())):
+        if not (identity := context.context_variables.get("identity", ContextVarsIdentity())):
             context.context_variables["identity"] = identity
 
         # Look for the api key value in the request. If it is not found, raise an exception
@@ -95,7 +96,7 @@ def validate_api_key(
                     return api_key if api_key["enabled"] else None
 
     def get_request_api_key(
-        self, rest_api: RestApi, request: InvocationRequest, identity: IdentityContext
+        self, rest_api: RestApi, request: InvocationRequest, identity: ContextVarsIdentity
     ) -> Optional[str]:
         """https://docs.aws.amazon.com/apigateway/latest/developerguide/api-gateway-api-key-source.html
         The source of the API key for metering requests according to a usage plan.
diff --git a/tests/unit/services/apigateway/test_handler_api_key_validation.py b/tests/unit/services/apigateway/test_handler_api_key_validation.py
@@ -7,13 +7,15 @@
 from localstack.services.apigateway.models import MergedRestApi, RestApiDeployment
 from localstack.services.apigateway.next_gen.execute_api.api import RestApiGatewayHandlerChain
 from localstack.services.apigateway.next_gen.execute_api.context import (
-    ContextVariables,
-    IdentityContext,
     InvocationRequest,
     RestApiInvocationContext,
 )
 from localstack.services.apigateway.next_gen.execute_api.gateway_response import InvalidAPIKeyError
 from localstack.services.apigateway.next_gen.execute_api.handlers import ApiKeyValidationHandler
+from localstack.services.apigateway.next_gen.execute_api.variables import (
+    ContextVariables,
+    ContextVarsIdentity,
+)
 from localstack.testing.config import TEST_AWS_ACCOUNT_ID, TEST_AWS_REGION_NAME
 
 TEST_API_ID = "testapi"
@@ -71,7 +73,7 @@ def _create_context(
 
         # Context populated by a Lambda Authorizer
         if api_key is not None:
-            context.context_variables["identity"] = IdentityContext(apiKey=api_key)
+            context.context_variables["identity"] = ContextVarsIdentity(apiKey=api_key)
         return context
 
     return _create_context
