feat: create BaseInfrastructureV2 to demonstrate before/after

Signed-off-by: heitorlessa <lessa@amazon.co.uk>

Author: heitorlessa

tests/e2e/metrics/infrastructure.py

@@ -1,126 +1,11 @@
 from pathlib import Path
-from typing import TYPE_CHECKING, Dict, Tuple
-from uuid import uuid4
 
-import boto3
-import yaml
-from aws_cdk import App, CfnOutput, RemovalPolicy, Stack, aws_logs
-from aws_cdk.aws_lambda import Code, Function, LayerVersion, Runtime, Tracing
+from tests.e2e.utils.infrastructure import BaseInfrastructureV2
 
-from tests.e2e.utils.asset import Assets
 
-if TYPE_CHECKING:
-    from mypy_boto3_cloudformation import CloudFormationClient
+class MetricsStack(BaseInfrastructureV2):
+    def __init__(self, handlers_dir: Path, feature_name: str = "metrics") -> None:
+        super().__init__(feature_name, handlers_dir)
 
-
-class MetricsStack:
-    def __init__(self, handlers_dir: Path) -> None:
-        self.stack_name = f"test-metrics-{uuid4()}"
-        self.handlers_dir = handlers_dir
-        self.app = App()
-        self.stack = Stack(self.app, self.stack_name)
-        self.session = boto3.Session()
-        self.cf_client: "CloudFormationClient" = self.session.client("cloudformation")
-        # NOTE: CDK stack account and region are tokens, we need to resolve earlier
-        self.account_id = self.session.client("sts").get_caller_identity()["Account"]
-        self.region = self.session.region_name
-        self.stack_outputs: Dict[str, str] = {}
-
-    def create_functions(self):
-        handlers = list(self.handlers_dir.rglob("*.py"))
-        source = Code.from_asset(f"{self.handlers_dir}")
-        for fn in handlers:
-            fn_name = fn.stem
-            function_python = Function(
-                self.stack,
-                id=f"{fn_name}-lambda",
-                runtime=Runtime.PYTHON_3_8,
-                code=source,
-                handler=f"{fn_name}.lambda_handler",
-                tracing=Tracing.ACTIVE,
-                layers=[
-                    LayerVersion.from_layer_version_arn(
-                        self.stack,
-                        f"{fn_name}-lambda-powertools",
-                        f"arn:aws:lambda:{self.region}:017000801446:layer:AWSLambdaPowertoolsPython:29",
-                    )
-                ],
-            )
-
-            aws_logs.LogGroup(
-                self.stack,
-                id=f"{fn_name}-lg",
-                log_group_name=f"/aws/lambda/{function_python.function_name}",
-                retention=aws_logs.RetentionDays.ONE_DAY,
-                removal_policy=RemovalPolicy.DESTROY,
-            )
-
-            # CFN Outputs only support hyphen
-            fn_name_camel_case = fn_name.title().replace("_", "")  # basic_handler -> BasicHandler
-            self._add_resource_output(
-                name=fn_name_camel_case, value=function_python.function_name, arn_value=function_python.function_arn
-            )
-
-    def deploy(self) -> Dict[str, str]:
-        """Creates CloudFormation Stack and return stack outputs as dict
-
-        Returns
-        -------
-        Dict[str, str]
-            CloudFormation Stack Outputs with output key and value
-        """
-        template, asset_manifest_file = self._synthesize()
-        assets = Assets(cfn_template=asset_manifest_file, account_id=self.account_id, region=self.region)
-        assets.upload()
-        return self._deploy_stack(self.stack_name, template)
-
-    def delete(self):
-        self.cf_client.delete_stack(StackName=self.stack_name)
-
-    def get_stack_outputs(self) -> Dict[str, str]:
-        return self.stack_outputs
-
-    def _synthesize(self) -> Tuple[Dict, Path]:
-        self.create_functions()
-        cloud_assembly = self.app.synth()
-        cf_template: Dict = cloud_assembly.get_stack_by_name(self.stack_name).template
-        cloud_assembly_assets_manifest_path: str = (
-            cloud_assembly.get_stack_by_name(self.stack_name).dependencies[0].file
-        )
-        return cf_template, Path(cloud_assembly_assets_manifest_path)
-
-    def _deploy_stack(self, stack_name: str, template: Dict) -> Dict[str, str]:
-        self.cf_client.create_stack(
-            StackName=stack_name,
-            TemplateBody=yaml.dump(template),
-            TimeoutInMinutes=10,
-            OnFailure="ROLLBACK",
-            Capabilities=["CAPABILITY_IAM"],
-        )
-        waiter = self.cf_client.get_waiter("stack_create_complete")
-        waiter.wait(StackName=stack_name, WaiterConfig={"Delay": 10, "MaxAttempts": 50})
-
-        stack_details = self.cf_client.describe_stacks(StackName=stack_name)
-        stack_outputs = stack_details["Stacks"][0]["Outputs"]
-        self.stack_outputs = {
-            output["OutputKey"]: output["OutputValue"] for output in stack_outputs if output["OutputKey"]
-        }
-
-        return self.stack_outputs
-
-    def _add_resource_output(self, name: str, value: str, arn_value: str):
-        """Add both resource value and ARN as Outputs to facilitate tests.
-
-        This will create two outputs: {Name} and {Name}Arn
-
-        Parameters
-        ----------
-        name : str
-            CloudFormation Output Key
-        value : str
-            CloudFormation Output Value
-        arn_value : str
-            CloudFormation Output Value for ARN
-        """
-        CfnOutput(self.stack, f"{name}", value=value)
-        CfnOutput(self.stack, f"{name}Arn", value=arn_value)
+    def create_resources(self):
+        self.create_lambda_functions()

tests/e2e/metrics/test_metrics.py

@@ -67,6 +67,6 @@ def test_cold_start_metric(cold_start_fn_arn: str, cold_start_fn: str):
     assert metric_data and metric_data[0] == 1.0
 
 
+# Abstract Infrastructure fixture parallelization work
 # helpers: adjust retries and wait to be much smaller
 # helpers: make retry config adjustable
-# Create separate Infra class so they can live side by side

tests/e2e/utils/helpers.py

@@ -1,6 +1,6 @@
 import json
 import secrets
-from datetime import datetime
+from datetime import datetime, timedelta
 from functools import lru_cache
 from typing import Dict, List, Optional, Tuple, Union
 
@@ -108,7 +108,7 @@ def get_metrics(
         When no metric is found within retry window
     """
     cw_client = cw_client or boto3.client("cloudwatch")
-    end_date = end_date or start_date + datetime.timedelta(minutes=2)
+    end_date = end_date or start_date + timedelta(minutes=2)
 
     metric_query = build_metric_query_data(
         namespace=namespace, metric_name=metric_name, period=period, stat=stat, dimensions=dimensions

tests/e2e/utils/infrastructure.py

@@ -6,12 +6,16 @@
 from abc import ABC, abstractmethod
 from enum import Enum
 from pathlib import Path
-from typing import Dict, List, Tuple, Type
+from typing import Dict, List, Optional, Tuple, Type
+from uuid import uuid4
 
 import boto3
 import yaml
 from aws_cdk import App, AssetStaging, BundlingOptions, CfnOutput, DockerImage, RemovalPolicy, Stack, aws_logs
 from aws_cdk.aws_lambda import Code, Function, LayerVersion, Runtime, Tracing
+from mypy_boto3_cloudformation import CloudFormationClient
+
+from tests.e2e.utils.asset import Assets
 
 PYTHON_RUNTIME_VERSION = f"V{''.join(map(str, sys.version_info[:2]))}"
 
@@ -24,11 +28,11 @@ class PythonVersion(Enum):
 
 class BaseInfrastructureStack(ABC):
     @abstractmethod
-    def synthesize() -> Tuple[dict, str]:
+    def synthesize(self) -> Tuple[dict, str]:
         ...
 
     @abstractmethod
-    def __call__() -> Tuple[dict, str]:
+    def __call__(self) -> Tuple[dict, str]:
         ...
 
 
@@ -210,3 +214,174 @@ def _find_assets(self, asset_template: str, account_id: str, region: str):
 
     def _transform_output(self, outputs: dict):
         return {output["OutputKey"]: output["OutputValue"] for output in outputs if output["OutputKey"]}
+
+
+class BaseInfrastructureV2(ABC):
+    def __init__(self, feature_name: str, handlers_dir: Path) -> None:
+        self.stack_name = f"test-{feature_name}-{uuid4()}"
+        self.handlers_dir = handlers_dir
+        self.app = App()
+        self.stack = Stack(self.app, self.stack_name)
+        self.session = boto3.Session()
+        self.cf_client: CloudFormationClient = self.session.client("cloudformation")
+        # NOTE: CDK stack account and region are tokens, we need to resolve earlier
+        self.account_id = self.session.client("sts").get_caller_identity()["Account"]
+        self.region = self.session.region_name
+        self.stack_outputs: Dict[str, str] = {}
+
+    def create_lambda_functions(self, function_props: Optional[Dict] = None):
+        """Create Lambda functions available under handlers_dir
+
+        It creates CloudFormation Outputs for every function found in PascalCase. For example,
+        {handlers_dir}/basic_handler.py creates `BasicHandler` and `BasicHandlerArn` outputs.
+
+
+        Parameters
+        ----------
+        function_props: Optional[Dict]
+            CDK Lambda FunctionProps as dictionary to override defaults
+
+        Examples
+        -------
+
+        Creating Lambda functions available in the handlers directory
+
+        ```python
+        self.create_lambda_functions()
+        ```
+
+        Creating Lambda functions and override runtime to Python 3.7
+
+        ```python
+        from aws_cdk.aws_lambda import Runtime
+
+        self.create_lambda_functions(function_props={"runtime": Runtime.PYTHON_3_7)
+        ```
+        """
+        handlers = list(self.handlers_dir.rglob("*.py"))
+        source = Code.from_asset(f"{self.handlers_dir}")
+        props_override = function_props or {}
+
+        for fn in handlers:
+            fn_name = fn.stem
+            function_settings = {
+                "id": f"{fn_name}-lambda",
+                "code": source,
+                "handler": f"{fn_name}.lambda_handler",
+                "tracing": Tracing.ACTIVE,
+                "runtime": Runtime.PYTHON_3_9,
+                "layers": [
+                    LayerVersion.from_layer_version_arn(
+                        self.stack,
+                        f"{fn_name}-lambda-powertools",
+                        f"arn:aws:lambda:{self.region}:017000801446:layer:AWSLambdaPowertoolsPython:29",
+                    )
+                ],
+                **props_override,
+            }
+
+            function_python = Function(self.stack, **function_settings)
+
+            aws_logs.LogGroup(
+                self.stack,
+                id=f"{fn_name}-lg",
+                log_group_name=f"/aws/lambda/{function_python.function_name}",
+                retention=aws_logs.RetentionDays.ONE_DAY,
+                removal_policy=RemovalPolicy.DESTROY,
+            )
+
+            # CFN Outputs only support hyphen
+            fn_name_pascal_case = fn_name.title().replace("_", "")  # basic_handler -> BasicHandler
+            self._add_resource_output(
+                name=fn_name_pascal_case, value=function_python.function_name, arn=function_python.function_arn
+            )
+
+    def deploy(self) -> Dict[str, str]:
+        """Creates CloudFormation Stack and return stack outputs as dict
+
+        Returns
+        -------
+        Dict[str, str]
+            CloudFormation Stack Outputs with output key and value
+        """
+        template, asset_manifest_file = self._synthesize()
+        assets = Assets(cfn_template=asset_manifest_file, account_id=self.account_id, region=self.region)
+        assets.upload()
+        return self._deploy_stack(self.stack_name, template)
+
+    def delete(self):
+        self.cf_client.delete_stack(StackName=self.stack_name)
+
+    def get_stack_outputs(self) -> Dict[str, str]:
+        return self.stack_outputs
+
+    @abstractmethod
+    def create_resources(self):
+        """Create any necessary CDK resources. It'll be called before deploy
+
+        Examples
+        -------
+
+        Creating a S3 bucket and export name and ARN
+
+        ```python
+        def created_resources(self):
+            s3 = s3.Bucket(self.stack, "MyBucket")
+
+            # This will create MyBucket and MyBucketArn CloudFormation Output
+            self._add_resource_output(name="MyBucket", value=s3.bucket_name, arn_value=bucket.bucket_arn)
+        ```
+
+        Creating Lambda functions available in the handlers directory
+
+        ```python
+        def created_resources(self):
+            self.create_lambda_functions()
+        ```
+        """
+        ...
+
+    def _synthesize(self) -> Tuple[Dict, Path]:
+        self.create_resources()
+        cloud_assembly = self.app.synth()
+        cf_template: Dict = cloud_assembly.get_stack_by_name(self.stack_name).template
+        cloud_assembly_assets_manifest_path: str = (
+            cloud_assembly.get_stack_by_name(self.stack_name).dependencies[0].file  # type: ignore[attr-defined]
+        )
+        return cf_template, Path(cloud_assembly_assets_manifest_path)
+
+    def _deploy_stack(self, stack_name: str, template: Dict) -> Dict[str, str]:
+        self.cf_client.create_stack(
+            StackName=stack_name,
+            TemplateBody=yaml.dump(template),
+            TimeoutInMinutes=10,
+            OnFailure="ROLLBACK",
+            Capabilities=["CAPABILITY_IAM"],
+        )
+        waiter = self.cf_client.get_waiter("stack_create_complete")
+        waiter.wait(StackName=stack_name, WaiterConfig={"Delay": 10, "MaxAttempts": 50})
+
+        stack_details = self.cf_client.describe_stacks(StackName=stack_name)
+        stack_outputs = stack_details["Stacks"][0]["Outputs"]
+        self.stack_outputs = {
+            output["OutputKey"]: output["OutputValue"] for output in stack_outputs if output["OutputKey"]
+        }
+
+        return self.stack_outputs
+
+    def _add_resource_output(self, name: str, value: str, arn: str):
+        """Add both resource value and ARN as Outputs to facilitate tests.
+
+        This will create two outputs: {Name} and {Name}Arn
+
+        Parameters
+        ----------
+        name : str
+            CloudFormation Output Key
+        value : str
+            CloudFormation Output Value
+        arn : str
+            CloudFormation Output Value for ARN
+        """
+        CfnOutput(self.stack, f"{name}", value=value)
+        CfnOutput(self.stack, f"{name}Arn", value=arn)