adding lambda functions

Author: rpmcginty

src/aibs_informatics_aws_utils/lambda_.py

@@ -0,0 +1,46 @@
+from typing import TYPE_CHECKING, List, Optional, Union
+
+from aibs_informatics_core.models.aws.core import AWSAccountId, AWSRegion
+from aibs_informatics_core.models.aws.lambda_ import LambdaFunctionName
+from botocore.exceptions import ClientError
+
+from aibs_informatics_aws_utils.core import AWSService, get_client_error_code
+
+if TYPE_CHECKING:  # pragma: no cover
+    from mypy_boto3_lambda.type_defs import FileSystemConfigTypeDef
+else:
+    FileSystemConfigTypeDef = dict
+
+
+get_lambda_client = AWSService.LAMBDA.get_client
+
+
+def get_lambda_function_url(
+    function_name: Union[LambdaFunctionName, str], region: AWSRegion = None
+) -> Optional[str]:
+    function_name = LambdaFunctionName(function_name)
+
+    lambda_client = get_lambda_client(region=region)
+
+    try:
+        response = lambda_client.get_function_url_config(FunctionName=function_name)
+    except ClientError as e:
+        if get_client_error_code(e) == "ResourceNotFoundException":
+            return None
+        else:
+            raise e
+    return response["FunctionUrl"]
+
+
+def get_lambda_function_file_systems(
+    function_name: Union[LambdaFunctionName, str], region: AWSRegion = None
+) -> List[FileSystemConfigTypeDef]:
+    function_name = LambdaFunctionName(function_name)
+
+    lambda_client = get_lambda_client(region=region)
+
+    response = lambda_client.get_function_configuration(FunctionName=function_name)
+
+    fs_configs = response.get("FileSystemConfigs")
+
+    return fs_configs

test/aibs_informatics_aws_utils/test_lambda.py

@@ -0,0 +1,116 @@
+import json
+from test.aibs_informatics_aws_utils.base import AwsBaseTest
+
+import boto3
+from moto import mock_iam, mock_lambda_simple, mock_sts
+from pytest import raises
+
+from aibs_informatics_aws_utils.lambda_ import (
+    get_lambda_function_file_systems,
+    get_lambda_function_url,
+)
+
+
+@mock_sts
+class LambdaTests(AwsBaseTest):
+    def setUp(self) -> None:
+        super().setUp()
+        self.set_aws_credentials()
+
+    def get_role_arn(self) -> str:
+        return boto3.client("iam").create_role(
+            RoleName="foo",
+            AssumeRolePolicyDocument=json.dumps(
+                {
+                    "Version": "2012-10-17",
+                    "Statement": [
+                        {
+                            "Effect": "Allow",
+                            "Principal": {"Service": "lambda.amazonaws.com"},
+                            "Action": "sts:AssumeRole",
+                        }
+                    ],
+                }
+            ),
+        )["Role"]["Arn"]
+
+    @mock_iam
+    @mock_lambda_simple
+    def test__get_lambda_function_file_systems__no_file_systems(self):
+        # Set up lambda
+        lambda_client = boto3.client("lambda")
+        file_system_configs = [
+            {
+                "Arn": "arn:aws:elasticfilesystem:us-west-2:123456789012:access-point/fsap-1234abcd",
+                "LocalMountPath": "/mnt/efs1",
+            },
+            {
+                "Arn": "arn:aws:elasticfilesystem:us-west-2:123456789012:access-point/fsap-1234abcd2",
+                "LocalMountPath": "/mnt/efs2",
+            },
+        ]
+        lambda_client.create_function(
+            FunctionName="test",
+            Runtime="python3.8",
+            Handler="test",
+            Role=self.get_role_arn(),
+            Code={"ZipFile": b"bar"},
+            Environment={"Variables": {"TEST": "test"}},
+            # NOTE: FileSystemConfigs is not supported by moto yet, so this is meaningless
+            FileSystemConfigs=file_system_configs,
+        )
+        # HACK: moto doesn't support FileSystemConfigs yet, so we have to patch it in
+        #       here we will fetch the actual response and then add the FileSystemConfigs
+        response = lambda_client.get_function_configuration(FunctionName="test")
+        with self.stub(lambda_client) as lambda_stubber:
+            lambda_stubber.add_response(
+                "get_function_configuration",
+                {
+                    **response,
+                    **{"FileSystemConfigs": file_system_configs},
+                },
+                expected_params={"FunctionName": "test"},
+            )
+            self.create_patch(
+                "aibs_informatics_aws_utils.lambda_.get_lambda_client", return_value=lambda_client
+            )
+
+            actual_file_system_configs = get_lambda_function_file_systems("test")
+            self.assertListEqual(actual_file_system_configs, file_system_configs)
+
+    @mock_iam
+    @mock_lambda_simple
+    def test__get_lambda_function_url__with_url(self):
+        lambda_client = boto3.client("lambda")
+        lambda_client.create_function(
+            FunctionName="test",
+            Runtime="python3.8",
+            Handler="test",
+            Role=self.get_role_arn(),
+            Code={"ZipFile": b"bar"},
+            Environment={"Variables": {"TEST": "test"}},
+        )
+        response = lambda_client.create_function_url_config(
+            FunctionName="test", AuthType="AWS_IAM"
+        )
+
+        assert get_lambda_function_url("test") == response["FunctionUrl"]
+
+    @mock_iam
+    @mock_lambda_simple
+    def test__get_lambda_function_url__handles_no_url(self):
+        lambda_client = boto3.client("lambda")
+        lambda_client.create_function(
+            FunctionName="test",
+            Runtime="python3.8",
+            Handler="test",
+            Role=self.get_role_arn(),
+            Code={"ZipFile": b"bar"},
+            Environment={"Variables": {"TEST": "test"}},
+        )
+
+    @mock_iam
+    @mock_lambda_simple
+    def test__get_lambda_function_url__handles_invalid_function_name(self):
+        with raises(ValueError):
+            get_lambda_function_url("@#$#$@#$@#$")