PR suggestions applied

Author: Michal Ploski

tests/e2e/conftest.py

@@ -30,8 +30,12 @@ class InfrastructureOutput:
     def get_lambda_arns(self) -> Dict[str, str]:
         return self.arns
 
-    def get_lambda_arn(self, name: str) -> Optional[str]:
-        return self.arns.get(name)
+    def get_lambda_function_arn(self, cf_output_name: str) -> Optional[str]:
+        return self.arns.get(cf_output_name)
+
+    def get_lambda_function_name(self, cf_output_name: str) -> Optional[str]:
+        lambda_arn = self.get_lambda_function_arn(cf_output_name=cf_output_name)
+        return lambda_arn.split(":")[-1] if lambda_arn else None
 
     def get_lambda_execution_time(self) -> datetime.datetime:
         return self.execution_time

tests/e2e/logger/test_logger.py

@@ -18,14 +18,12 @@ def config() -> conftest.LambdaConfig:
 
 def test_basic_lambda_logs_visible(execute_lambda: conftest.InfrastructureOutput, config: conftest.LambdaConfig):
     # GIVEN
-    lambda_arn = execute_lambda.get_lambda_arn(name="basichandlerarn")
+    lambda_name = execute_lambda.get_lambda_function_name(cf_output_name="basichandlerarn")
     timestamp = execute_lambda.get_lambda_execution_time_timestamp()
     cw_client = boto3.client("logs")
 
     # WHEN
-    filtered_logs = helpers.get_logs(
-        lambda_function_name=lambda_arn.split(":")[-1], start_time=timestamp, log_client=cw_client
-    )
+    filtered_logs = helpers.get_logs(lambda_function_name=lambda_name, start_time=timestamp, log_client=cw_client)
 
     # THEN
     assert any(
@@ -39,14 +37,12 @@ def test_basic_lambda_no_debug_logs_visible(
     execute_lambda: conftest.InfrastructureOutput, config: conftest.LambdaConfig
 ):
     # GIVEN
-    lambda_arn = execute_lambda.get_lambda_arn(name="basichandlerarn")
+    lambda_name = execute_lambda.get_lambda_function_name(cf_output_name="basichandlerarn")
     timestamp = execute_lambda.get_lambda_execution_time_timestamp()
     cw_client = boto3.client("logs")
 
     # WHEN
-    filtered_logs = helpers.get_logs(
-        lambda_function_name=lambda_arn.split(":")[-1], start_time=timestamp, log_client=cw_client
-    )
+    filtered_logs = helpers.get_logs(lambda_function_name=lambda_name, start_time=timestamp, log_client=cw_client)
 
     # THEN
     assert not any(
@@ -65,14 +61,12 @@ def test_basic_lambda_contextual_data_logged(execute_lambda: conftest.Infrastruc
         "cold_start",
     )
 
-    lambda_arn = execute_lambda.get_lambda_arn(name="basichandlerarn")
+    lambda_name = execute_lambda.get_lambda_function_name(cf_output_name="basichandlerarn")
     timestamp = execute_lambda.get_lambda_execution_time_timestamp()
     cw_client = boto3.client("logs")
 
     # WHEN
-    filtered_logs = helpers.get_logs(
-        lambda_function_name=lambda_arn.split(":")[-1], start_time=timestamp, log_client=cw_client
-    )
+    filtered_logs = helpers.get_logs(lambda_function_name=lambda_name, start_time=timestamp, log_client=cw_client)
 
     # THEN
     assert all(keys in logs.dict(exclude_unset=True) for logs in filtered_logs for keys in required_keys)
@@ -82,14 +76,12 @@ def test_basic_lambda_additional_key_persistence_basic_lambda(
     execute_lambda: conftest.InfrastructureOutput, config: conftest.LambdaConfig
 ):
     # GIVEN
-    lambda_arn = execute_lambda.get_lambda_arn(name="basichandlerarn")
+    lambda_name = execute_lambda.get_lambda_function_name(cf_output_name="basichandlerarn")
     timestamp = execute_lambda.get_lambda_execution_time_timestamp()
     cw_client = boto3.client("logs")
 
     # WHEN
-    filtered_logs = helpers.get_logs(
-        lambda_function_name=lambda_arn.split(":")[-1], start_time=timestamp, log_client=cw_client
-    )
+    filtered_logs = helpers.get_logs(lambda_function_name=lambda_name, start_time=timestamp, log_client=cw_client)
 
     # THEN
     assert any(
@@ -103,14 +95,12 @@ def test_basic_lambda_additional_key_persistence_basic_lambda(
 def test_basic_lambda_empty_event_logged(execute_lambda: conftest.InfrastructureOutput):
 
     # GIVEN
-    lambda_arn = execute_lambda.get_lambda_arn(name="basichandlerarn")
+    lambda_name = execute_lambda.get_lambda_function_name(cf_output_name="basichandlerarn")
     timestamp = execute_lambda.get_lambda_execution_time_timestamp()
     cw_client = boto3.client("logs")
 
     # WHEN
-    filtered_logs = helpers.get_logs(
-        lambda_function_name=lambda_arn.split(":")[-1], start_time=timestamp, log_client=cw_client
-    )
+    filtered_logs = helpers.get_logs(lambda_function_name=lambda_name, start_time=timestamp, log_client=cw_client)
 
     # THEN
     assert any(log.message == {} for log in filtered_logs)
@@ -127,14 +117,12 @@ def test_no_context_lambda_contextual_data_not_logged(execute_lambda: conftest.I
         "cold_start",
     )
 
-    lambda_arn = execute_lambda.get_lambda_arn(name="nocontexthandlerarn")
+    lambda_name = execute_lambda.get_lambda_function_name(cf_output_name="nocontexthandlerarn")
     timestamp = execute_lambda.get_lambda_execution_time_timestamp()
     cw_client = boto3.client("logs")
 
     # WHEN
-    filtered_logs = helpers.get_logs(
-        lambda_function_name=lambda_arn.split(":")[-1], start_time=timestamp, log_client=cw_client
-    )
+    filtered_logs = helpers.get_logs(lambda_function_name=lambda_name, start_time=timestamp, log_client=cw_client)
 
     # THEN
     assert not any(keys in logs.dict(exclude_unset=True) for logs in filtered_logs for keys in required_missing_keys)
@@ -143,14 +131,12 @@ def test_no_context_lambda_contextual_data_not_logged(execute_lambda: conftest.I
 def test_no_context_lambda_event_not_logged(execute_lambda: conftest.InfrastructureOutput):
 
     # GIVEN
-    lambda_arn = execute_lambda.get_lambda_arn(name="nocontexthandlerarn")
+    lambda_name = execute_lambda.get_lambda_function_name(cf_output_name="nocontexthandlerarn")
     timestamp = execute_lambda.get_lambda_execution_time_timestamp()
     cw_client = boto3.client("logs")
 
     # WHEN
-    filtered_logs = helpers.get_logs(
-        lambda_function_name=lambda_arn.split(":")[-1], start_time=timestamp, log_client=cw_client
-    )
+    filtered_logs = helpers.get_logs(lambda_function_name=lambda_name, start_time=timestamp, log_client=cw_client)
 
     # THEN
     assert not any(log.message == {} for log in filtered_logs)

tests/e2e/metrics/test_metrics.py

@@ -3,9 +3,8 @@
 
 import boto3
 import pytest
-
-from .. import conftest
-from ..utils import helpers
+from e2e import conftest
+from e2e.utils import helpers
 
 
 @pytest.fixture(scope="module")
@@ -15,7 +14,7 @@ def config() -> conftest.LambdaConfig:
         "environment_variables": {
             "POWERTOOLS_METRICS_NAMESPACE": "powertools-e2e-metric",
             "POWERTOOLS_SERVICE_NAME": "test-powertools-service",
-            "METRIC_NAME": f"business-metric-{uuid.uuid4()}",
+            "METRIC_NAME": f"business-metric-{str(uuid.uuid4()).replace('-','_')}",
         },
     }
 

tests/e2e/tracer/handlers/basic_handler.py

@@ -13,11 +13,13 @@
 
 @tracer.capture_lambda_handler
 def lambda_handler(event: dict, context: LambdaContext):
+    tracer.put_annotation(key=ANNOTATION_KEY, value=ANNOTATION_VALUE)
     tracer.put_metadata(key=ANNOTATION_KEY, value=ANNOTATION_VALUE)
     return asyncio.run(collect_payment())
 
 
 @tracer.capture_method
 async def collect_payment() -> str:
+    tracer.put_annotation(key=ANNOTATION_KEY, value=ANNOTATION_ASYNC_VALUE)
     tracer.put_metadata(key=ANNOTATION_KEY, value=ANNOTATION_ASYNC_VALUE)
     return "success"

tests/e2e/tracer/test_tracer.py

@@ -8,11 +8,11 @@
 
 
 @pytest.fixture(scope="module")
-def config():
+def config() -> conftest.LambdaConfig:
     return {
         "parameters": {"tracing": "ACTIVE"},
         "environment_variables": {
-            "ANNOTATION_KEY": f"e2e-tracer-{uuid.uuid4()}",
+            "ANNOTATION_KEY": f"e2e-tracer-{str(uuid.uuid4()).replace('-','_')}",
             "ANNOTATION_VALUE": "stored",
             "ANNOTATION_ASYNC_VALUE": "payments",
         },
@@ -21,10 +21,10 @@ def config():
 
 def test_basic_lambda_async_trace_visible(execute_lambda: conftest.InfrastructureOutput, config: conftest.LambdaConfig):
     # GIVEN
-    lambda_arn = execute_lambda.get_lambda_arn(name="basichandlerarn")
+    lambda_name = execute_lambda.get_lambda_function_name(cf_output_name="basichandlerarn")
     start_date = execute_lambda.get_lambda_execution_time()
     end_date = start_date + datetime.timedelta(minutes=5)
-    trace_filter_exporession = f'service("{lambda_arn.split(":")[-1]}")'
+    trace_filter_exporession = f'service("{lambda_name}")'
 
     # WHEN
     trace = helpers.get_traces(

tests/e2e/utils/helpers.py

@@ -107,6 +107,7 @@ def get_traces(filter_expression: str, xray_client: XRayClient, start_date: date
 
 
 def find_trace_additional_info(trace: Dict) -> List[TraceSegment]:
+    """Find all trace annotations and metadata and return them to the caller"""
     info = []
     for segment in trace["Traces"][0]["Segments"]:
         document = json.loads(segment["Document"])