-
Notifications
You must be signed in to change notification settings - Fork 14
Add proxy support and tests for Lambda service #128
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
whummer
wants to merge
1
commit into
main
Choose a base branch
from
add-lambda-proxy-tests
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
+298
−1
Open
Changes from all commits
Commits
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,280 @@ | ||
| # 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" |
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I'm not quite sure if we should consider the invoking of lambdas as a "read" operation.
Running a lambda can have many side-effects. I wouldn't expect the lambda to be called when I use read-only mode. Especially in the era of AI-agents, this can be very problematic for a user that may have thought that they give read-only access to the proxy.
Could we introduce some other mode specifically for cases like this? (like
executeorinvoke?)