APIGW: Fix tests to run with non-default account ID (localstack#9690)

matt-mercer · Dec 11, 2023 · 215fa1b · 215fa1b
1 parent a4f0a22
commit 215fa1b
Show file tree

Hide file tree

Showing 5 changed files with 141 additions and 45 deletions.
diff --git a/localstack/services/apigateway/helpers.py b/localstack/services/apigateway/helpers.py
@@ -585,8 +585,10 @@ def get_stage_variables(context: ApiInvocationContext) -> Optional[Dict[str, str
     if not context.stage:
         return {}
 
-    _, region_name = get_api_account_id_and_region(context.api_id)
-    api_gateway_client = connect_to(region_name=region_name).apigateway
+    account_id, region_name = get_api_account_id_and_region(context.api_id)
+    api_gateway_client = connect_to(
+        aws_access_key_id=account_id, region_name=region_name
+    ).apigateway
     try:
         response = api_gateway_client.get_stage(restApiId=context.api_id, stageName=context.stage)
         return response.get("variables")

diff --git a/localstack/services/apigateway/integration.py b/localstack/services/apigateway/integration.py
@@ -168,9 +168,9 @@ def get_source_arn(invocation_context: ApiInvocationContext):
 def call_lambda(
     function_arn: str, event: bytes, asynchronous: bool, invocation_context: ApiInvocationContext
 ) -> str:
-    region_name = extract_region_from_arn(function_arn)
     clients = get_service_factory(
-        region_name=region_name, role_arn=invocation_context.integration.get("credentials")
+        region_name=extract_region_from_arn(function_arn),
+        role_arn=invocation_context.integration.get("credentials"),
     )
     inv_result = clients.lambda_.request_metadata(
         service_principal=ServicePrincipal.apigateway, source_arn=get_source_arn(invocation_context)
@@ -754,7 +754,11 @@ def invoke(self, invocation_context: ApiInvocationContext):
         else:
             payload = json.loads(invocation_context.data)
 
-        client = connect_to().stepfunctions
+        client = get_service_factory(
+            region_name=invocation_context.region_name,
+            role_arn=invocation_context.integration.get("credentials"),
+        ).stepfunctions
+
         if isinstance(payload.get("input"), dict):
             payload["input"] = json.dumps(payload["input"])
 

diff --git a/localstack/utils/testutil.py b/localstack/utils/testutil.py
@@ -290,6 +290,7 @@ def connect_api_gateway_to_http_with_lambda_proxy(
     auth_creator_func=None,
     http_method=None,
     client=None,
+    role_arn: str = None,
 ):
     if methods is None:
         methods = []
@@ -304,12 +305,15 @@ def connect_api_gateway_to_http_with_lambda_proxy(
 
     for method in methods:
         int_meth = http_method or method
+        integration = {"type": "AWS_PROXY", "uri": target_uri, "httpMethod": int_meth}
+        if role_arn:
+            integration["credentials"] = role_arn
         resources[resource_path].append(
             {
                 "httpMethod": method,
                 "authorizationType": auth_type,
                 "authorizerId": None,
-                "integrations": [{"type": "AWS_PROXY", "uri": target_uri, "httpMethod": int_meth}],
+                "integrations": [integration],
             }
         )
     return resource_utils.create_api_gateway(
@@ -332,6 +336,7 @@ def create_lambda_api_gateway_integration(
     stage_name=None,
     auth_type=None,
     auth_creator_func=None,
+    role_arn: str = None,
 ):
     if methods is None:
         methods = []
@@ -355,6 +360,7 @@ def create_lambda_api_gateway_integration(
         methods=methods,
         auth_type=auth_type,
         auth_creator_func=auth_creator_func,
+        role_arn=role_arn,
     )
     return result
 

diff --git a/tests/aws/services/apigateway/conftest.py b/tests/aws/services/apigateway/conftest.py
@@ -43,6 +43,11 @@
     "Statement": [{"Effect": "Allow", "Action": "lambda:*", "Resource": "*"}],
 }
 
+APIGATEWAY_S3_POLICY = {
+    "Version": "2012-10-17",
+    "Statement": [{"Effect": "Allow", "Action": "s3:*", "Resource": "*"}],
+}
+
 APIGATEWAY_DYNAMODB_POLICY = {
     "Version": "2012-10-17",
     "Statement": [{"Effect": "Allow", "Action": "dynamodb:*", "Resource": "*"}],