test_lambda.py

# Note: This file has been (partially or fully) generated by an AI agent.
import io
import json
import logging
import zipfile

import boto3
import pytest
from botocore.exceptions import ClientError
from localstack.aws.connect import connect_to
from localstack.utils.strings import short_uid
from localstack.utils.sync import retry

from aws_proxy.shared.models import ProxyConfig

LOG = logging.getLogger(__name__)

LAMBDA_RUNTIME = "python3.12"
LAMBDA_HANDLER = "index.handler"
LAMBDA_HANDLER_CODE = """
import json
def handler(event, context):
    return {"statusCode": 200, "body": json.dumps({"message": "hello", "input": event})}
"""


def _create_lambda_zip() -> bytes:
    """Create an in-memory ZIP deployment package with a simple handler."""
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
        zf.writestr("index.py", LAMBDA_HANDLER_CODE)
    return buf.getvalue()


@pytest.fixture(scope="module")
def lambda_execution_role():
    """Create a basic IAM role for Lambda execution in real AWS, shared across tests in this module."""
    iam_client = boto3.client("iam")
    role_name = f"test-lambda-proxy-role-{short_uid()}"
    trust_policy = json.dumps(
        {
            "Version": "2012-10-17",
            "Statement": [
                {
                    "Effect": "Allow",
                    "Principal": {"Service": "lambda.amazonaws.com"},
                    "Action": "sts:AssumeRole",
                }
            ],
        }
    )
    role = iam_client.create_role(
        RoleName=role_name,
        AssumeRolePolicyDocument=trust_policy,
        Description="Test role for Lambda proxy tests",
    )
    role_arn = role["Role"]["Arn"]

    # Wait for the role to be usable by Lambda (IAM eventual consistency)
    def _wait_for_role():
        iam_client.get_role(RoleName=role_name)

    retry(_wait_for_role, retries=10, sleep=2)

    yield role_arn

    # cleanup
    try:
        iam_client.delete_role(RoleName=role_name)
    except Exception as e:
        LOG.warning("Failed to clean up IAM role %s: %s", role_name, e)


def _create_lambda_function(
    lambda_client_aws, function_name, role_arn, cleanups, env_vars=None
):
    """Helper to create a Lambda function in real AWS and register cleanup."""
    kwargs = {
        "FunctionName": function_name,
        "Runtime": LAMBDA_RUNTIME,
        "Role": role_arn,
        "Handler": LAMBDA_HANDLER,
        "Code": {"ZipFile": _create_lambda_zip()},
        "Timeout": 30,
    }
    if env_vars:
        kwargs["Environment"] = {"Variables": env_vars}

    # Retry create_function to handle IAM role propagation delay
    def _create():
        lambda_client_aws.create_function(**kwargs)

    retry(_create, retries=15, sleep=5)
    cleanups.append(
        lambda: lambda_client_aws.delete_function(FunctionName=function_name)
    )

    # Wait for function to become Active
    def _wait_active():
        config = lambda_client_aws.get_function_configuration(
            FunctionName=function_name
        )
        if config["State"] != "Active":
            raise AssertionError(
                f"Function {function_name} not active yet: {config['State']}"
            )

    retry(_wait_active, retries=30, sleep=2)


def test_lambda_requests(start_aws_proxy, cleanups, lambda_execution_role):
    """Test basic Lambda proxy: create in AWS, describe/invoke via LocalStack proxy."""
    function_name_aws = f"test-fn-aws-{short_uid()}"
    function_name_local = f"test-fn-local-{short_uid()}"

    # start proxy - only forwarding requests for function name matching the AWS function
    config = ProxyConfig(
        services={"lambda": {"resources": f".*:function:{function_name_aws}"}}
    )
    start_aws_proxy(config)

    # create clients
    region_name = "us-east-1"
    lambda_client = connect_to(region_name=region_name).lambda_
    lambda_client_aws = boto3.client("lambda", region_name=region_name)

    # create function in real AWS
    _create_lambda_function(
        lambda_client_aws, function_name_aws, lambda_execution_role, cleanups
    )

    # assert that local call for GetFunction is proxied and returns AWS data
    fn_local = lambda_client.get_function(FunctionName=function_name_aws)
    fn_aws = lambda_client_aws.get_function(FunctionName=function_name_aws)
    assert fn_local["Configuration"]["FunctionName"] == function_name_aws
    assert (
        fn_local["Configuration"]["FunctionArn"]
        == fn_aws["Configuration"]["FunctionArn"]
    )
    assert fn_local["Configuration"]["Runtime"] == LAMBDA_RUNTIME

    # assert that GetFunctionConfiguration is also proxied
    config_local = lambda_client.get_function_configuration(
        FunctionName=function_name_aws
    )
    config_aws = lambda_client_aws.get_function_configuration(
        FunctionName=function_name_aws
    )
    assert config_local["FunctionArn"] == config_aws["FunctionArn"]
    assert config_local["Handler"] == LAMBDA_HANDLER

    # invoke function through proxy and verify it executes on real AWS
    response_local = lambda_client.invoke(
        FunctionName=function_name_aws,
        Payload=json.dumps({"key": "value"}),
    )
    payload_local = json.loads(response_local["Payload"].read())
    assert payload_local["statusCode"] == 200
    body = json.loads(payload_local["body"])
    assert body["message"] == "hello"
    assert body["input"] == {"key": "value"}

    # invoke via AWS client directly and compare
    response_aws = lambda_client_aws.invoke(
        FunctionName=function_name_aws,
        Payload=json.dumps({"key": "value"}),
    )
    payload_aws = json.loads(response_aws["Payload"].read())
    assert payload_aws["statusCode"] == 200

    # negative test: a non-matching function should NOT exist in AWS
    with pytest.raises(ClientError) as ctx:
        lambda_client_aws.get_function(FunctionName=function_name_local)
    assert ctx.value.response["Error"]["Code"] == "ResourceNotFoundException"


def test_lambda_read_only(start_aws_proxy, cleanups, lambda_execution_role):
    """Test Lambda proxy in read-only mode: reads proxied, writes/invokes blocked."""
    function_name = f"test-fn-ro-{short_uid()}"

    # start proxy in read-only mode with wildcard resources
    config = ProxyConfig(services={"lambda": {"resources": ".*", "read_only": True}})
    start_aws_proxy(config)

    # create clients
    region_name = "us-east-1"
    lambda_client = connect_to(region_name=region_name).lambda_
    lambda_client_aws = boto3.client("lambda", region_name=region_name)

    # create function in real AWS (direct, not through proxy)
    _create_lambda_function(
        lambda_client_aws, function_name, lambda_execution_role, cleanups
    )

    # read operations should be proxied
    fn_local = lambda_client.get_function(FunctionName=function_name)
    fn_aws = lambda_client_aws.get_function(FunctionName=function_name)
    assert (
        fn_local["Configuration"]["FunctionArn"]
        == fn_aws["Configuration"]["FunctionArn"]
    )

    config_local = lambda_client.get_function_configuration(FunctionName=function_name)
    assert config_local["FunctionArn"] == fn_aws["Configuration"]["FunctionArn"]

    # ListFunctions should also be proxied (read operation)
    def _list_contains_function():
        functions = lambda_client.list_functions()["Functions"]
        func_names = [f["FunctionName"] for f in functions]
        assert function_name in func_names

    retry(_list_contains_function, retries=5, sleep=2)

    # Invoke is classified as a read operation for proxy purposes
    # (it executes the function but doesn't modify it)
    response = lambda_client.invoke(
        FunctionName=function_name,
        Payload=json.dumps({"key": "value"}),
    )
    payload = json.loads(response["Payload"].read())
    assert payload["statusCode"] == 200
    body = json.loads(payload["body"])
    assert body["message"] == "hello"

    # UpdateFunctionConfiguration is a write operation - should be blocked
    with pytest.raises(ClientError) as ctx:
        lambda_client.update_function_configuration(
            FunctionName=function_name,
            Description="updated via proxy - should not work",
        )
    assert ctx.value.response["Error"]["Code"] == "ResourceNotFoundException"


def test_lambda_resource_name_matching(
    start_aws_proxy, cleanups, lambda_execution_role
):
    """Test that only functions matching the resource pattern are proxied."""
    fn_match = f"proxy-fn-{short_uid()}"
    fn_nomatch = f"local-fn-{short_uid()}"

    # start proxy - only forwarding requests for functions starting with "proxy-fn-"
    config = ProxyConfig(services={"lambda": {"resources": ".*:function:proxy-fn-.*"}})
    start_aws_proxy(config)

    # create clients
    region_name = "us-east-1"
    lambda_client = connect_to(region_name=region_name).lambda_
    lambda_client_aws = boto3.client("lambda", region_name=region_name)

    # create matching function in real AWS
    _create_lambda_function(
        lambda_client_aws, fn_match, lambda_execution_role, cleanups
    )

    # matching function should be accessible through proxy
    fn_local = lambda_client.get_function(FunctionName=fn_match)
    fn_aws = lambda_client_aws.get_function(FunctionName=fn_match)
    assert (
        fn_local["Configuration"]["FunctionArn"]
        == fn_aws["Configuration"]["FunctionArn"]
    )

    # invoke matching function through proxy
    response = lambda_client.invoke(
        FunctionName=fn_match,
        Payload=json.dumps({"test": True}),
    )
    payload = json.loads(response["Payload"].read())
    assert payload["statusCode"] == 200

    # non-matching function should NOT exist in AWS (negative test)
    with pytest.raises(ClientError) as ctx:
        lambda_client_aws.get_function(FunctionName=fn_nomatch)
    assert ctx.value.response["Error"]["Code"] == "ResourceNotFoundException"

    # GetFunction for non-matching name through local client should NOT be proxied
    # (goes to LocalStack which doesn't have it either)
    with pytest.raises(ClientError) as ctx:
        lambda_client.get_function(FunctionName=fn_nomatch)
    assert ctx.value.response["Error"]["Code"] == "ResourceNotFoundException"