109 lines
3.4 KiB
Python
109 lines
3.4 KiB
Python
import json
|
|
import os
|
|
from functools import wraps
|
|
from time import sleep
|
|
from uuid import uuid4
|
|
|
|
import boto3
|
|
from botocore.exceptions import ClientError
|
|
|
|
from moto import mock_aws
|
|
|
|
|
|
def lambda_aws_verified(func):
|
|
"""
|
|
Function that is verified to work against AWS.
|
|
Can be run against AWS at any time by setting:
|
|
MOTO_TEST_ALLOW_AWS_REQUEST=true
|
|
|
|
If this environment variable is not set, the function runs in a `mock_aws` context.
|
|
|
|
This decorator will:
|
|
- Create an IAM-role that can be used by AWSLambda functions table
|
|
- Run the test
|
|
- Delete the role
|
|
"""
|
|
|
|
@wraps(func)
|
|
def pagination_wrapper():
|
|
role_name = "moto_test_role_" + str(uuid4())[0:6]
|
|
|
|
allow_aws_request = (
|
|
os.environ.get("MOTO_TEST_ALLOW_AWS_REQUEST", "false").lower() == "true"
|
|
)
|
|
|
|
if allow_aws_request:
|
|
return create_role_and_test(role_name)
|
|
else:
|
|
with mock_aws():
|
|
return create_role_and_test(role_name)
|
|
|
|
def create_role_and_test(role_name):
|
|
from .test_lambda import _lambda_region
|
|
|
|
policy_doc = {
|
|
"Version": "2012-10-17",
|
|
"Statement": [
|
|
{
|
|
"Sid": "LambdaAssumeRole",
|
|
"Effect": "Allow",
|
|
"Principal": {"Service": "lambda.amazonaws.com"},
|
|
"Action": "sts:AssumeRole",
|
|
}
|
|
],
|
|
}
|
|
iam = boto3.client("iam", region_name=_lambda_region)
|
|
iam_role_arn = iam.create_role(
|
|
RoleName=role_name,
|
|
AssumeRolePolicyDocument=json.dumps(policy_doc),
|
|
Path="/",
|
|
)["Role"]["Arn"]
|
|
|
|
try:
|
|
# Role is not immediately available
|
|
check_iam_role_is_ready_to_use(iam_role_arn, role_name, _lambda_region)
|
|
|
|
resp = func(iam_role_arn)
|
|
finally:
|
|
iam.delete_role(RoleName=role_name)
|
|
|
|
return resp
|
|
|
|
def check_iam_role_is_ready_to_use(role_arn: str, role_name: str, region: str):
|
|
from .utilities import get_test_zip_file1
|
|
|
|
iam = boto3.client("iam", region_name=region)
|
|
|
|
# The waiter is normally used to wait until a resource is ready
|
|
wait_for_role = iam.get_waiter("role_exists")
|
|
wait_for_role.wait(RoleName=role_name)
|
|
|
|
# HOWEVER:
|
|
# Just because the role exists, doesn't mean it can be used by Lambda
|
|
# The only reliable way to check if our role is ready:
|
|
# - try to create a function and see if it succeeds
|
|
#
|
|
# The alternive is to wait between 5 and 10 seconds, which is not a great solution either
|
|
_lambda = boto3.client("lambda", region)
|
|
for _ in range(15):
|
|
try:
|
|
fn_name = f"fn_for_{role_name}"
|
|
_lambda.create_function(
|
|
FunctionName=fn_name,
|
|
Runtime="python3.11",
|
|
Role=role_arn,
|
|
Handler="lambda_function.lambda_handler",
|
|
Code={"ZipFile": get_test_zip_file1()},
|
|
)
|
|
# If it succeeded, our role is ready
|
|
_lambda.delete_function(FunctionName=fn_name)
|
|
|
|
return
|
|
except ClientError:
|
|
sleep(1)
|
|
raise Exception(
|
|
f"Couldn't create test Lambda FN using IAM role {role_name}, probably because it wasn't ready in time"
|
|
)
|
|
|
|
return pagination_wrapper
|