From ebd22f875b6c82cb97047849973070733573b44b Mon Sep 17 00:00:00 2001 From: Gerald Manipon Date: Thu, 23 May 2024 10:37:25 -0700 Subject: [PATCH] add capability to configure router via ROUTER_CFG_URL env variable --- README.md | 4 +- pyproject.toml | 4 ++ scripts/build_mock_lambda_package.sh | 1 + src/unity_initiator/cloud/lambda_handler.py | 19 +++++-- tests/test_lambda.py | 57 +++++++++++++-------- 5 files changed, 59 insertions(+), 26 deletions(-) diff --git a/README.md b/README.md index ba863e1..00fd463 100644 --- a/README.md +++ b/README.md @@ -110,10 +110,10 @@ This guide provides a quick way to get started with our project. Please see our dist ├── unity_initiator-0.0.1-py3-none-any.whl └── unity_initiator-0.0.1.tar.gz - + 1 directory, 2 files ``` - + ### Test Instructions diff --git a/pyproject.toml b/pyproject.toml index 2f80db2..b0dd278 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,7 @@ dependencies = [ "httpx~=0.27.0", "respx~=0.21.1", "docker~=7.0.0", + "smart-open~=7.0.4", # TODO: remove cryptography pin when this issue resolved: # cryptography/hazmat/bindings/_rust.abi3.so: cannot open shared object file: No such file or directory "cryptography<35.0.0", @@ -100,3 +101,6 @@ exclude_lines = [ "if __name__ == .__main__.:", "if TYPE_CHECKING:", ] + +[tool.isort] +profile = "black" \ No newline at end of file diff --git a/scripts/build_mock_lambda_package.sh b/scripts/build_mock_lambda_package.sh index 65748b1..70edad4 100755 --- a/scripts/build_mock_lambda_package.sh +++ b/scripts/build_mock_lambda_package.sh @@ -15,5 +15,6 @@ VERSION=$(hatch run python -c 'from importlib.metadata import version; print(ver mkdir -p $PKG_DIR pip install -t $PKG_DIR ${DIST_DIR}/unity_initiator-*.whl cp ${TEST_DIR}/test_lambda.py $PKG_DIR/lambda_function.py +cp -r ${TEST_DIR} $PKG_DIR/ cd $PKG_DIR zip -rq ${DIST_DIR}/unity_initiator-${VERSION}-mock_lambda.zip . diff --git a/src/unity_initiator/cloud/lambda_handler.py b/src/unity_initiator/cloud/lambda_handler.py index 9192c08..6e4c39e 100644 --- a/src/unity_initiator/cloud/lambda_handler.py +++ b/src/unity_initiator/cloud/lambda_handler.py @@ -3,6 +3,8 @@ from html import unescape from tempfile import mkstemp +import smart_open + from ..router import Router from ..utils.logger import logger @@ -11,14 +13,23 @@ def lambda_handler_base(event, context): """Base lambda handler that instantiates a router, globally, and executes actions for a single payload.""" + logger.info("context: %s", context) - # TODO: Should use either AppConfig or retrieve router config from S3 location. - # For now, reading router config body from ROUTER_CFG env variable then writing - # to local file. + # TODO: Should use AppConfig. For now, either reading router config body in ROUTER_CFG env variable + # or from a url in ROUTER_CFG_URL env variable. global ROUTER if ROUTER is None: - router_cfg = os.environ["ROUTER_CFG"] + router_cfg = os.environ.get("ROUTER_CFG", "").strip() + router_cfg_url = os.environ.get("ROUTER_CFG_URL", "").strip() + if router_cfg == "": + if router_cfg_url != "": + with smart_open.open(router_cfg_url, "r") as f: + router_cfg = f.read() + else: + raise RuntimeError( + "No router configuration specified via ROUTER_CFG or ROUTER_CFG_URL env variables." + ) fd, router_file = mkstemp(prefix="router_", suffix=".yaml", text=True) with os.fdopen(fd, "w") as f: f.write(router_cfg) diff --git a/tests/test_lambda.py b/tests/test_lambda.py index b4f0a88..fc0be00 100644 --- a/tests/test_lambda.py +++ b/tests/test_lambda.py @@ -8,6 +8,7 @@ import docker import pytest import respx +import smart_open from botocore.exceptions import ClientError from httpx import Response from importlib_resources import files @@ -24,14 +25,34 @@ os.environ["AWS_DEFAULT_REGION"] = "hilo-hawaii-1" +# test bucket for mock +TEST_BUCKET = "test_bucket" + + def setup_mock_resources(): """Create mocked AWS and Airflow resources.""" + # create router config in mock S3 bucket + # TODO: Should use AppConfig. For now, writing out router config to + # an S3 url to pass in as the ROUTER_CFG_URL env variable. + router_file = files("tests.resources").joinpath("test_router.yaml") + with open(router_file) as f: + router_cfg = f.read() + s3_client = boto3.client("s3") + s3_client.create_bucket( + Bucket=TEST_BUCKET, + CreateBucketConfiguration={ + "LocationConstraint": os.environ["AWS_DEFAULT_REGION"] + }, + ) + with smart_open.open(f"s3://{TEST_BUCKET}/test_router.yaml", "w") as f: + f.write(router_cfg) + # create mock SNS topics - client = boto3.client("sns") - client.create_topic(Name="eval_sbg_l2_readiness") - client.create_topic(Name="eval_m2020_xyz_left_finder") - client.create_topic(Name="eval_nisar_ingest") + sns_client = boto3.client("sns") + sns_client.create_topic(Name="eval_sbg_l2_readiness") + sns_client.create_topic(Name="eval_m2020_xyz_left_finder") + sns_client.create_topic(Name="eval_nisar_ingest") # mock airflow REST API respx.post("https://example.com/api/v1/dags/eval_nisar_l0a_readiness/dagRuns").mock( @@ -136,8 +157,8 @@ def setup_class(cls): cls.client = boto3.client("lambda") cls.function_name = str(uuid4())[0:6] - # TODO: Should use either AppConfig or retrieve router config from S3 location. - # For now, writing out router config body to env variable to pass to lambda. + # TODO: Should use AppConfig. For now, writing out router config body to + # the ROUTER_CFG env variable. cls.router_file = files("tests.resources").joinpath("test_router.yaml") with open(cls.router_file) as f: cls.router_cfg = f.read() @@ -257,7 +278,7 @@ def setup_class(cls): cls.logs_client = boto3.client("logs") cls.s3_client = boto3.client("s3") - cls.bucket_name = "test_bucket" + cls.bucket_name = TEST_BUCKET cls.bucket = cls.s3_client.create_bucket( Bucket=cls.bucket_name, CreateBucketConfiguration={ @@ -304,8 +325,7 @@ def setup_class(cls): QueueUrl=cls.sqs_queue_initiator["QueueUrl"], MaxNumberOfMessages=10 ) assert len(messages["Messages"]) == 1 - # print("SQS message:") - # print(json.dumps(messages["Messages"], indent=2)) + logger.debug("SQS message: %s", json.dumps(messages["Messages"], indent=2)) cls.sqs_client.delete_message( QueueUrl=cls.sqs_queue_initiator["QueueUrl"], ReceiptHandle=messages["Messages"][0]["ReceiptHandle"], @@ -313,21 +333,14 @@ def setup_class(cls): message_body = messages["Messages"][0]["Body"] sns_message = json.loads(message_body) assert sns_message["Type"] == "Notification" - # print("SNS message:") - # print(json.dumps(sns_message, indent=2)) + logger.debug("SNS message: %s", json.dumps(sns_message, indent=2)) # get S3 notification from SNS message s3_message_body = json.loads(sns_message["Message"]) assert s3_message_body["Event"] == "s3:TestEvent" - # print("S3 message:") - # print(json.dumps(s3_message_body, indent=2)) - - # TODO: Should use either AppConfig or retrieve router config from S3 location. - # For now, writing out router config body to env variable to pass to lambda. - cls.router_file = files("tests.resources").joinpath("test_router.yaml") - with open(cls.router_file) as f: - cls.router_cfg = f.read() + logger.debug("S3 message: %s", json.dumps(s3_message_body, indent=2)) + # create mocked initiator lambda cls.fxn = cls.lambda_client.create_function( FunctionName=cls.function_name, Runtime="python3.11", @@ -338,7 +351,11 @@ def setup_class(cls): Timeout=3, MemorySize=128, Publish=True, - Environment={"Variables": {"ROUTER_CFG": cls.router_cfg}}, + Environment={ + "Variables": { + "ROUTER_CFG_URL": f"s3://{cls.bucket_name}/test_router.yaml" + } + }, ) # create event source mapping