moto/tests/test_awslambda/__init__.py

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