From 90d1bacb58addac100130394118e4402b6b314a1 Mon Sep 17 00:00:00 2001
From: Pierre Trespeuch <trespeuch@adacore.com>
Date: Mon, 2 Oct 2023 11:48:13 +0200
Subject: [PATCH] Add a method to add outputs to a Stack

This is convenient to export values that can be imported by other
CloudFormation stacks.
---
 src/e3/aws/troposphere/__init__.py            |  9 ++-
 .../troposphere/stack/stack_test.py           | 32 ++++++++
 .../troposphere/stack/stack_with_outputs.json | 78 +++++++++++++++++++
 3 files changed, 118 insertions(+), 1 deletion(-)
 create mode 100644 tests/tests_e3_aws/troposphere/stack/stack_with_outputs.json

diff --git a/src/e3/aws/troposphere/__init__.py b/src/e3/aws/troposphere/__init__.py
index fc5b614f..ea717df1 100644
--- a/src/e3/aws/troposphere/__init__.py
+++ b/src/e3/aws/troposphere/__init__.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 from abc import ABC, abstractmethod
 from itertools import chain
-from troposphere import AWSObject, Template
+from troposphere import AWSObject, Output, Template
 
 from e3.aws import cfn, name_to_id, Session
 from e3.aws.cfn.main import CFNMain
@@ -130,6 +130,13 @@ def add(self, element: AWSObject | Construct | Stack) -> Stack:
 
         return self
 
+    def add_output(self, output: Output | list[Output]) -> None:
+        """Add outputs to stack template.
+
+        :param output: output to add to the template
+        """
+        self.template.add_output(output)
+
     def add_condition(self, condition_name: str, condition: ConditionFunction) -> None:
         """Add condition to stack template.
 
diff --git a/tests/tests_e3_aws/troposphere/stack/stack_test.py b/tests/tests_e3_aws/troposphere/stack/stack_test.py
index da549bc4..3dc36746 100644
--- a/tests/tests_e3_aws/troposphere/stack/stack_test.py
+++ b/tests/tests_e3_aws/troposphere/stack/stack_test.py
@@ -1,8 +1,14 @@
 """Provide Stack tests."""
+import json
+from pathlib import Path
+
+from troposphere import Output, Export
 
 from e3.aws.troposphere.s3.bucket import Bucket
 from e3.aws.troposphere import Stack
 
+TEST_DIR = Path(__file__).parent
+
 
 def test_instanciate() -> None:
     """Test stack instanciation."""
@@ -16,3 +22,29 @@ def test_add_and_get_item() -> None:
     stack.add(Bucket("my-bucket"))
     my_bucket = stack["my-bucket"]
     assert my_bucket
+
+
+def test_add_outputs() -> None:
+    """Test adding outputs to a stack."""
+    stack = Stack("test-stack", "this is a test stack")
+    stack.add(Bucket("my-bucket"))
+    stack.add_output(
+        Output("MyOutput1", Description="My first output", Value=Export(name="Output1"))
+    )
+    stack.add_output(
+        [
+            Output(
+                "MyOutput2",
+                Description="My second output",
+                Value=Export(name="Output2"),
+            ),
+            Output(
+                "MyOutput3", Description="My third output", Value=Export(name="Output3")
+            ),
+        ]
+    )
+
+    with open(TEST_DIR / "stack_with_outputs.json") as fd:
+        expected_template = json.load(fd)
+
+    assert stack.export()["Resources"] == expected_template
diff --git a/tests/tests_e3_aws/troposphere/stack/stack_with_outputs.json b/tests/tests_e3_aws/troposphere/stack/stack_with_outputs.json
new file mode 100644
index 00000000..a567f7e4
--- /dev/null
+++ b/tests/tests_e3_aws/troposphere/stack/stack_with_outputs.json
@@ -0,0 +1,78 @@
+{
+    "MyBucket": {
+        "Properties": {
+            "BucketName": "my-bucket",
+            "BucketEncryption": {
+                "ServerSideEncryptionConfiguration": [
+                    {
+                        "ServerSideEncryptionByDefault": {
+                            "SSEAlgorithm": "AES256"
+                        }
+                    }
+                ]
+            },
+            "PublicAccessBlockConfiguration": {
+                "BlockPublicAcls": true,
+                "BlockPublicPolicy": true,
+                "IgnorePublicAcls": true,
+                "RestrictPublicBuckets": true
+            },
+            "VersioningConfiguration": {
+                "Status": "Enabled"
+            }
+        },
+        "Type": "AWS::S3::Bucket"
+    },
+    "MyBucketPolicy": {
+        "Properties": {
+            "Bucket": {
+                "Ref": "MyBucket"
+            },
+            "PolicyDocument": {
+                "Version": "2012-10-17",
+                "Statement": [
+                    {
+                        "Effect": "Deny",
+                        "Principal": {
+                            "AWS": "*"
+                        },
+                        "Action": "s3:*",
+                        "Resource": "arn:aws:s3:::my-bucket/*",
+                        "Condition": {
+                            "Bool": {
+                                "aws:SecureTransport": "false"
+                            }
+                        }
+                    },
+                    {
+                        "Effect": "Deny",
+                        "Principal": {
+                            "AWS": "*"
+                        },
+                        "Action": "s3:PutObject",
+                        "Resource": "arn:aws:s3:::my-bucket/*",
+                        "Condition": {
+                            "StringNotEquals": {
+                                "s3:x-amz-server-side-encryption": "AES256"
+                            }
+                        }
+                    },
+                    {
+                        "Effect": "Deny",
+                        "Principal": {
+                            "AWS": "*"
+                        },
+                        "Action": "s3:PutObject",
+                        "Resource": "arn:aws:s3:::my-bucket/*",
+                        "Condition": {
+                            "Null": {
+                                "s3:x-amz-server-side-encryption": "true"
+                            }
+                        }
+                    }
+                ]
+            }
+        },
+        "Type": "AWS::S3::BucketPolicy"
+    }
+}
\ No newline at end of file