refactor(helpers): create builders for data and types to simplify writing test

Author: heitorlessa

tests/e2e/metrics/conftest.py

@@ -1,13 +1,32 @@
+from pathlib import Path
+
 import pytest
+from _pytest import fixtures
 
 from tests.e2e.metrics.infrastructure import MetricsStack
 
 
 @pytest.fixture(autouse=True)
-def infrastructure() -> MetricsStack:
-    # Use request fixture to remove hardcode handler dir
+def infrastructure(request: fixtures.SubRequest) -> MetricsStack:
+    """Setup and teardown logic for E2E test infrastructure
+
+    Parameters
+    ----------
+    request : fixtures.SubRequest
+        test fixture containing metadata about test execution
+
+    Returns
+    -------
+    MetricsStack
+        Metrics Stack to deploy infrastructure
+
+    Yields
+    ------
+    Iterator[MetricsStack]
+        Deployed Infrastructure
+    """
     try:
-        stack = MetricsStack()
+        stack = MetricsStack(handlers_dir=Path(f"{request.fspath.dirname}/handlers"))
         stack.deploy()
         yield stack
     finally:

tests/e2e/metrics/infrastructure.py

@@ -14,9 +14,9 @@
 
 
 class MetricsStack:
-    def __init__(self) -> None:
+    def __init__(self, handlers_dir: Path) -> None:
         self.stack_name = f"test-metrics-{uuid4()}"
-        self.handlers_dir = "tests/e2e/metrics/handlers"  # hardcoded as we're not using a fixture yet
+        self.handlers_dir = handlers_dir  # hardcoded as we're not using a fixture yet
         self.app = App()
         self.stack = Stack(self.app, self.stack_name)
         self.session = boto3.Session()
@@ -27,8 +27,8 @@ def __init__(self) -> None:
         self.stack_outputs: Dict[str, str] = {}
 
     def create_functions(self):
-        handlers = list(Path(self.handlers_dir).rglob("*.py"))
-        source = Code.from_asset(self.handlers_dir)
+        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(
@@ -54,7 +54,10 @@ def create_functions(self):
                 retention=aws_logs.RetentionDays.ONE_DAY,
                 removal_policy=RemovalPolicy.DESTROY,
             )
-            CfnOutput(self.stack, f"{fn_name}_arn", value=function_python.function_arn)
+
+            # CFN Outputs only support hyphen
+            fn_name_camel_case = fn_name.title().replace("_", "")  # basic_handler -> BasicHandler
+            CfnOutput(self.stack, f"{fn_name_camel_case}Arn", value=function_python.function_arn)
 
     def deploy(self) -> Dict[str, str]:
         """Creates CloudFormation Stack and return stack outputs as dict

tests/e2e/metrics/test_metrics.py

@@ -1,36 +1,37 @@
 import datetime
 import json
-from typing import Dict, Type
 
 import pytest
-from e2e.utils import helpers
 
 from tests.e2e.metrics.infrastructure import MetricsStack
+from tests.e2e.utils import helpers
+
+METRIC_NAMESPACE = "powertools-e2e-metric"
 
 
 @pytest.fixture
-def infra_outputs(infrastructure: Type[MetricsStack]):
+def infra_outputs(infrastructure: MetricsStack):
     return infrastructure.get_stack_outputs()
 
 
-def test_basic_lambda_metric_is_visible(infra_outputs: Dict[str, str]):
-    # GIVEN
-    metric_name = "test-two"
-    service = "test-metric-is-visible"
-    namespace = "powertools-e2e-metric"
-    event = json.dumps({"metric_name": metric_name, "service": service, "namespace": namespace})
+@pytest.fixture
+def basic_handler_fn(infra_outputs: dict) -> str:
+    return infra_outputs.get("BasicHandlerArn", "")
 
-    # NOTE: Need to try creating a dynamic enum/dataclass w/ Literal types to make discovery easier
-    # it might not be possible
+
+def test_basic_lambda_metric_is_visible(basic_handler_fn):
+    # GIVEN
+    metric_name = helpers.build_metric_name()
+    service = helpers.build_service_name()
+    event = json.dumps({"metric_name": metric_name, "service": service, "namespace": METRIC_NAMESPACE})
 
     # WHEN
-    execution_time = datetime.datetime.utcnow()
-    ret = helpers.trigger_lambda(lambda_arn=infra_outputs.get("basichandlerarn"), payload=event)
+    ret, execution_time = helpers.trigger_lambda(lambda_arn=basic_handler_fn, payload=event)
 
     metrics = helpers.get_metrics(
         start_date=execution_time,
         end_date=execution_time + datetime.timedelta(minutes=2),
-        namespace=namespace,
+        namespace=METRIC_NAMESPACE,
         service_name=service,
         metric_name=metric_name,
     )
@@ -48,8 +49,6 @@ def test_basic_lambda_metric_is_visible(infra_outputs: Dict[str, str]):
 
 # helpers: adjust retries and wait to be much smaller
 # helpers: make retry config adjustable
-# Infra: Create dynamic Enum/DataClass to reduce guessing on outputs
-# Infra: Fix outputs
 # Infra: Add temporary Powertools Layer
 # Powertools: should have a method to set namespace at runtime
 # Create separate Infra class so they can live side by side

tests/e2e/utils/helpers.py

@@ -1,15 +1,15 @@
 import json
+import secrets
 from datetime import datetime
 from functools import lru_cache
-from typing import TYPE_CHECKING, Dict, List, Optional, Union
-
-if TYPE_CHECKING:
-    from mypy_boto3_cloudwatch import MetricDataResultTypeDef
-    from mypy_boto3_cloudwatch.client import CloudWatchClient
-    from mypy_boto3_lambda.client import LambdaClient
-    from mypy_boto3_xray.client import XRayClient
+from typing import Dict, List, Optional, Tuple, Union
 
 import boto3
+from mypy_boto3_cloudwatch.client import CloudWatchClient
+from mypy_boto3_cloudwatch.type_defs import DimensionTypeDef, MetricDataQueryTypeDef, MetricDataResultTypeDef
+from mypy_boto3_lambda.client import LambdaClient
+from mypy_boto3_lambda.type_defs import InvocationResponseTypeDef
+from mypy_boto3_xray.client import XRayClient
 from pydantic import BaseModel
 from retry import retry
 
@@ -36,14 +36,17 @@ class TraceSegment(BaseModel):
     annotations: Dict = {}
 
 
-def trigger_lambda(lambda_arn: str, payload: str, client: Optional["LambdaClient"] = None):
+def trigger_lambda(
+    lambda_arn: str, payload: str, client: Optional[LambdaClient] = None
+) -> Tuple[InvocationResponseTypeDef, datetime]:
     client = client or boto3.client("lambda")
-    return client.invoke(FunctionName=lambda_arn, InvocationType="RequestResponse", Payload=payload)
+    execution_time = datetime.utcnow()
+    return client.invoke(FunctionName=lambda_arn, InvocationType="RequestResponse", Payload=payload), execution_time
 
 
 @lru_cache(maxsize=10, typed=False)
 @retry(ValueError, delay=1, jitter=1, tries=20)
-def get_logs(lambda_function_name: str, log_client: "CloudWatchClient", start_time: int, **kwargs: dict) -> List[Log]:
+def get_logs(lambda_function_name: str, log_client: CloudWatchClient, start_time: int, **kwargs: dict) -> List[Log]:
     response = log_client.filter_log_events(logGroupName=f"/aws/lambda/{lambda_function_name}", startTime=start_time)
     if not response["events"]:
         raise ValueError("Empty response from Cloudwatch Logs. Repeating...")
@@ -58,26 +61,24 @@ def get_logs(lambda_function_name: str, log_client: "CloudWatchClient", start_ti
     return filtered_logs
 
 
-@lru_cache(maxsize=10, typed=False)
-@retry(ValueError, delay=1, jitter=1, tries=5)
+@retry(ValueError, delay=1, jitter=1, tries=10)
 def get_metrics(
     namespace: str,
     start_date: datetime,
     metric_name: str,
     service_name: str,
-    cw_client: Optional["CloudWatchClient"] = None,
+    cw_client: Optional[CloudWatchClient] = None,
     end_date: Optional[datetime] = None,
-) -> "MetricDataResultTypeDef":
+    period: int = 60,
+    stat: str = "Sum",
+) -> MetricDataResultTypeDef:
     cw_client = cw_client or boto3.client("cloudwatch")
-    metric_query = {
-        "Id": "m1",
-        "Expression": f'SELECT MAX("{metric_name}") from SCHEMA("{namespace}",service) \
-                    where service=\'{service_name}\'',
-        "ReturnData": True,
-        "Period": 600,
-    }
+    metric_query = build_metric_query_data(
+        namespace=namespace, metric_name=metric_name, service_name=service_name, period=period, stat=stat
+    )
+
     response = cw_client.get_metric_data(
-        MetricDataQueries=[metric_query],
+        MetricDataQueries=metric_query,
         StartTime=start_date,
         EndTime=end_date or datetime.utcnow(),
     )
@@ -88,7 +89,10 @@ def get_metrics(
 
 
 @retry(ValueError, delay=1, jitter=1, tries=10)
-def get_traces(filter_expression: str, xray_client: "XRayClient", start_date: datetime, end_date: datetime) -> Dict:
+def get_traces(
+    filter_expression: str, start_date: datetime, end_date: datetime, xray_client: Optional[XRayClient] = None
+) -> Dict:
+    xray_client = xray_client or boto3.client("xray")
     paginator = xray_client.get_paginator("get_trace_summaries")
     response_iterator = paginator.paginate(
         StartTime=start_date,
@@ -132,3 +136,41 @@ def find_meta(segment: dict, result: List):
         )
         if x_subsegment.get("subsegments"):
             find_meta(segment=x_subsegment, result=result)
+
+
+# Maintenance: Build a separate module for builders
+def build_metric_name() -> str:
+    return f"test_metric{build_random_value()}"
+
+
+def build_service_name() -> str:
+    return f"test_service{build_random_value()}"
+
+
+def build_random_value(nbytes: int = 10) -> str:
+    return secrets.token_urlsafe(nbytes).replace("-", "")
+
+
+def build_metric_query_data(
+    namespace: str,
+    metric_name: str,
+    service_name: str,
+    period: int = 60,
+    stat: str = "Sum",
+    dimensions: Optional[DimensionTypeDef] = None,
+) -> List[MetricDataQueryTypeDef]:
+    metric_dimensions: List[DimensionTypeDef] = [{"Name": "service", "Value": service_name}]
+    if dimensions is not None:
+        metric_dimensions.append(dimensions)
+
+    return [
+        {
+            "Id": metric_name,
+            "MetricStat": {
+                "Metric": {"Namespace": namespace, "MetricName": metric_name, "Dimensions": metric_dimensions},
+                "Period": period,
+                "Stat": stat,
+            },
+            "ReturnData": True,
+        }
+    ]