Merge branch 'morosi-fix' into 'master'

Allow specify version or alias in add_lambda_subscription

See merge request it/e3-aws!88

Author: adanaja

src/e3/aws/troposphere/sns/__init__.py

@@ -11,7 +11,7 @@
     from typing import Any
     from troposphere import AWSObject
     from e3.aws.troposphere import Stack
-    from e3.aws.troposphere.awslambda import Function
+    from e3.aws.troposphere.awslambda import Function, Version, Alias
     from e3.aws.troposphere.iam.policy_statement import ConditionType
 
 
@@ -36,15 +36,19 @@ def _get_topic_policy_name(self) -> str:
         return name_to_id(f"{self.name}Policy")
 
     def add_lambda_subscription(
-        self, function: Function, delivery_policy: dict | None = None
+        self,
+        function: Function,
+        delivery_policy: dict | None = None,
+        version: Version | Alias | None = None,
     ) -> None:
         """Add a lambda subscription endpoint to topic.
 
         :param function: lambda function that will be added as endpoint
         :param delivery_policy: The delivery policy to assign to the subscription
+        :param version: specific version or alias to subscribe
         """
         sub_params = {
-            "Endpoint": function.arn,
+            "Endpoint": version.ref if version is not None else function.arn,
             "Protocol": "lambda",
             "TopicArn": self.arn,
         }
@@ -55,10 +59,18 @@ def add_lambda_subscription(
         self.optional_resources.extend(
             [
                 sns.SubscriptionResource(
-                    name_to_id(f"{function.name}Sub"), **sub_params
+                    name_to_id(
+                        "{}Sub".format(
+                            version.name if version is not None else function.name
+                        )
+                    ),
+                    **sub_params,
                 ),
                 function.invoke_permission(
-                    name_suffix=self.name, service="sns", source_arn=self.arn
+                    name_suffix=self.name,
+                    service="sns",
+                    source_arn=self.arn,
+                    version=version,
                 ),
             ]
         )

tests/tests_e3_aws/troposphere/sns/sns_test.py

@@ -1,9 +1,13 @@
 from __future__ import annotations
-
+from typing import TYPE_CHECKING
 import pytest
 
 from e3.aws.troposphere import Stack
 from e3.aws.troposphere.sns import Topic
+from e3.aws.troposphere.awslambda import PyFunction, Version, Alias
+
+if TYPE_CHECKING:
+    from typing import Any
 
 
 EXPECTED_TOPIC_DEFAULT_TEMPLATE = {
@@ -61,3 +65,86 @@ def test_allow_service_to_publish_not_unique_sid(stack: Stack) -> None:
         stack.add(topic)
 
     assert str(ex.value) == "Unique Sid is required for TopicPolicy statements"
+
+
+@pytest.mark.parametrize(
+    "version, expected_endpoint, expected_function_name_ref",
+    [
+        # Add the subscription to the function itself
+        (
+            None,
+            {"Fn::GetAtt": ["Mypylambda", "Arn"]},
+            "Mypylambda",
+        ),
+        # Add the subscription to a version of the function
+        (
+            Version(
+                name="myversion", description="this is some version", lambda_arn=""
+            ),
+            {"Ref": "Myversion"},
+            "Myversion",
+        ),
+        # Add the subscription to an alias of the function
+        (
+            Alias(
+                name="myalias",
+                description="this is some alias",
+                lambda_arn="",
+                lambda_version="",
+            ),
+            {"Ref": "Myalias"},
+            "Myalias",
+        ),
+    ],
+)
+def test_topic_lambda_subscription(
+    version: Version | Alias | None,
+    expected_endpoint: dict[str, Any],
+    expected_function_name_ref: str,
+    stack: Stack,
+) -> None:
+    """Test topic creation with lambda subscription.
+
+    :param version: a version or alias of the function
+    :param expected_endpoint: value that should be set for Endpoint
+    :param expected_function_name_ref: name that should be referenced in FunctionName
+    :param stack: the stack
+    """
+    topic = Topic("mytopic")
+    topic.add_lambda_subscription(
+        function=PyFunction(
+            name="mypylambda",
+            description="this is a test",
+            role="somearn",
+            runtime="python3.9",
+            code_dir="my_code_dir",
+            handler="app.main",
+        ),
+        version=version,
+    )
+    stack.add(topic)
+    assert stack.export()["Resources"] == EXPECTED_TOPIC_DEFAULT_TEMPLATE | {
+        f"{expected_function_name_ref}Sub": {
+            "Properties": {
+                "Endpoint": expected_endpoint,
+                "Protocol": "lambda",
+                "TopicArn": {
+                    "Ref": "Mytopic",
+                },
+            },
+            "Type": "AWS::SNS::Subscription",
+        },
+        f"{expected_function_name_ref}mytopic": {
+            "Properties": {
+                "Action": "lambda:InvokeFunction",
+                "FunctionName": {
+                    "Ref": expected_function_name_ref,
+                },
+                "Principal": "sns.amazonaws.com",
+                "SourceArn": {
+                    "Ref": "Mytopic",
+                },
+            },
+            "Type": "AWS::Lambda::Permission",
+        },
+    }