From f4490c226d1a8f903a25d47dc2c6fbd4ca548b13 Mon Sep 17 00:00:00 2001 From: Roy Moore Date: Wed, 15 May 2024 17:46:32 +0000 Subject: [PATCH] feat(core): Added AWS Lambda module --- index.rst | 1 + modules/aws/README.rst | 6 +++ modules/aws/testcontainers/aws/__init__.py | 1 + modules/aws/testcontainers/aws/aws_lambda.py | 52 +++++++++++++++++++ modules/aws/tests/lambda_sample/Dockerfile | 10 ++++ .../tests/lambda_sample/lambda_function.py | 5 ++ modules/aws/tests/test_aws.py | 38 ++++++++++++++ poetry.lock | 3 +- pyproject.toml | 2 + 9 files changed, 117 insertions(+), 1 deletion(-) create mode 100644 modules/aws/README.rst create mode 100644 modules/aws/testcontainers/aws/__init__.py create mode 100644 modules/aws/testcontainers/aws/aws_lambda.py create mode 100644 modules/aws/tests/lambda_sample/Dockerfile create mode 100644 modules/aws/tests/lambda_sample/lambda_function.py create mode 100644 modules/aws/tests/test_aws.py diff --git a/index.rst b/index.rst index eccaeef51..0c4209906 100644 --- a/index.rst +++ b/index.rst @@ -16,6 +16,7 @@ testcontainers-python facilitates the use of Docker containers for functional an core/README modules/arangodb/README + modules/awslambda/README modules/azurite/README modules/cassandra/README modules/chroma/README diff --git a/modules/aws/README.rst b/modules/aws/README.rst new file mode 100644 index 000000000..409e0ef5d --- /dev/null +++ b/modules/aws/README.rst @@ -0,0 +1,6 @@ +.. autoclass:: testcontainers.aws.AWSLambdaContainer +.. title:: testcontainers.aws.AWSLambdaContainer + +Make sure you are using an image based on `public.ecr.aws/lambda/python` + +Please checkout https://docs.aws.amazon.com/lambda/latest/dg/python-image.html for more information on how to run AWS Lambda functions locally. diff --git a/modules/aws/testcontainers/aws/__init__.py b/modules/aws/testcontainers/aws/__init__.py new file mode 100644 index 000000000..f16705c86 --- /dev/null +++ b/modules/aws/testcontainers/aws/__init__.py @@ -0,0 +1 @@ +from .aws_lambda import AWSLambdaContainer # noqa: F401 diff --git a/modules/aws/testcontainers/aws/aws_lambda.py b/modules/aws/testcontainers/aws/aws_lambda.py new file mode 100644 index 000000000..7bdab46ac --- /dev/null +++ b/modules/aws/testcontainers/aws/aws_lambda.py @@ -0,0 +1,52 @@ +from typing import Optional + +import httpx + +from testcontainers.core.generic import SrvContainer + +RIE_PATH = "/2015-03-31/functions/function/invocations" +# AWS OS-only base images contain an Amazon Linux distribution and the runtime interface emulator. + + +class AWSLambdaContainer(SrvContainer): + """ + AWS Lambda container that is based on a custom image. + + Example: + + .. doctest:: + + >>> from testcontainers.aws import AWSLambdaContainer + >>> from testcontainers.core.waiting_utils import wait_for_logs + + >>> with AWSLambdaContainer(path="./modules/aws/tests/lambda_sample", port=8080, tag="lambda_func:latest") as func: + ... response = func.send_request(data={'payload': 'some data'}) + ... assert response.status_code == 200 + ... assert "Hello from AWS Lambda using Python" in response.json() + ... delay = wait_for_logs(func, "START RequestId:") + """ + + def __init__(self, path: str, port: int, tag: Optional[str] = None, image_cleanup: bool = True) -> None: + """ + :param path: Path to the AWS Lambda dockerfile. + :param port: Port to expose the AWS Lambda function. + :param tag: Tag for the image to be built (default: None). + :param image_cleanup: Clean up the image after the container is stopped (default: True). + """ + super().__init__(path, port, tag, image_cleanup) + + def get_api_url(self) -> str: + return self._create_connection_url() + RIE_PATH + + def send_request(self, data: dict) -> httpx.Response: + """ + Send a request to the AWS Lambda function. + + :param data: Data to be sent to the AWS Lambda function. + :return: Response from the AWS Lambda function. + """ + client = httpx.Client() + return client.post(self.get_api_url(), json=data) + + def get_stdout(self) -> str: + return self.get_logs()[0].decode("utf-8") diff --git a/modules/aws/tests/lambda_sample/Dockerfile b/modules/aws/tests/lambda_sample/Dockerfile new file mode 100644 index 000000000..5d071c802 --- /dev/null +++ b/modules/aws/tests/lambda_sample/Dockerfile @@ -0,0 +1,10 @@ +FROM public.ecr.aws/lambda/python:3.9 + +RUN pip install boto3 + +COPY lambda_function.py ${LAMBDA_TASK_ROOT} + +EXPOSE 8080 + +# Set the CMD to your handler (could also be done as a parameter override outside of the Dockerfile) +CMD [ "lambda_function.handler" ] diff --git a/modules/aws/tests/lambda_sample/lambda_function.py b/modules/aws/tests/lambda_sample/lambda_function.py new file mode 100644 index 000000000..b253ed172 --- /dev/null +++ b/modules/aws/tests/lambda_sample/lambda_function.py @@ -0,0 +1,5 @@ +import sys + + +def handler(event, context): + return "Hello from AWS Lambda using Python" + sys.version + "!" diff --git a/modules/aws/tests/test_aws.py b/modules/aws/tests/test_aws.py new file mode 100644 index 000000000..f58b706f3 --- /dev/null +++ b/modules/aws/tests/test_aws.py @@ -0,0 +1,38 @@ +import re +import pytest + +from testcontainers.aws import AWSLambdaContainer +from testcontainers.aws.aws_lambda import RIE_PATH + +DOCKER_FILE_PATH = "./modules/aws/tests/lambda_sample" +IMAGE_TAG = "lambda:test" + + +def test_aws_lambda_container(): + with AWSLambdaContainer(path=DOCKER_FILE_PATH, port=8080, tag=IMAGE_TAG, image_cleanup=False) as func: + assert func.get_container_host_ip() == "localhost" + assert func.internal_port == 8080 + assert re.match(rf"http://localhost:\d+{RIE_PATH}", func.get_api_url()) + response = func.send_request(data={"payload": "test"}) + assert response.status_code == 200 + assert "Hello from AWS Lambda using Python" in response.json() + for log_str in ["START RequestId", "END RequestId", "REPORT RequestId"]: + assert log_str in func.get_stdout() + + +def test_aws_lambda_container_no_tag(): + with AWSLambdaContainer(path=DOCKER_FILE_PATH, port=8080, image_cleanup=True) as func: + response = func.send_request(data={"payload": "test"}) + assert response.status_code == 200 + + +def test_aws_lambda_container_no_port(): + with pytest.raises(TypeError): + with AWSLambdaContainer(path=DOCKER_FILE_PATH, tag=IMAGE_TAG, image_cleanup=False): + pass + + +def test_aws_lambda_container_no_path(): + with pytest.raises(TypeError): + with AWSLambdaContainer(port=8080, tag=IMAGE_TAG, image_cleanup=True): + pass diff --git a/poetry.lock b/poetry.lock index 7cb910630..7e689d26f 100644 --- a/poetry.lock +++ b/poetry.lock @@ -4178,6 +4178,7 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p [extras] arangodb = ["python-arango"] +aws = ["boto3"] azurite = ["azure-storage-blob"] cassandra = [] chroma = ["chromadb-client"] @@ -4213,4 +4214,4 @@ weaviate = ["weaviate-client"] [metadata] lock-version = "2.0" python-versions = ">=3.9,<4.0" -content-hash = "8676827f49237bd96b95c4cf9686728d401bddf78d09a2e3e018ac2dc882a8a9" +content-hash = "cc8b6fe17a69473f977d7819a893b87e6314f919f9aa9b27b8d2d05005a6a477" diff --git a/pyproject.toml b/pyproject.toml index 3d23423ab..b01aa3d4e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,6 +29,7 @@ classifiers = [ packages = [ { include = "testcontainers", from = "core" }, { include = "testcontainers", from = "modules/arangodb" }, + { include = "testcontainers", from = "modules/aws"}, { include = "testcontainers", from = "modules/azurite" }, { include = "testcontainers", from = "modules/cassandra" }, { include = "testcontainers", from = "modules/chroma" }, @@ -104,6 +105,7 @@ httpx = { version = "*", optional = true } [tool.poetry.extras] arangodb = ["python-arango"] +aws = ["boto3"] azurite = ["azure-storage-blob"] cassandra = [] clickhouse = ["clickhouse-driver"]