Create helper function for fetching trace metadata

Author: Michal Ploski

tests/e2e/logger/test_logger.py

@@ -60,7 +60,6 @@ def test_basic_lambda_no_debug_logs_visible(
 @pytest.mark.e2e
 def test_basic_lambda_contextual_data_logged(execute_lambda: conftest.InfrastructureOutput):
     # GIVEN
-
     required_keys = (
         "xray_trace_id",
         "function_request_id",
@@ -127,7 +126,6 @@ def test_basic_lambda_empty_event_logged(execute_lambda: conftest.Infrastructure
 def test_no_context_lambda_contextual_data_not_logged(execute_lambda: conftest.InfrastructureOutput):
 
     # GIVEN
-
     required_missing_keys = (
         "function_request_id",
         "function_arn",

tests/e2e/metrics/handlers/basic_handler.py

@@ -3,11 +3,9 @@
 from aws_lambda_powertools import Metrics
 from aws_lambda_powertools.metrics import MetricUnit
 
-METRIC_NAMESPACE = os.environ["METRIC_NAMESPACE"]
 METRIC_NAME = os.environ["METRIC_NAME"]
-SERVICE_NAME = os.environ["SERVICE_NAME"]
 
-metrics = Metrics(namespace=METRIC_NAMESPACE, service=SERVICE_NAME)
+metrics = Metrics()
 
 
 @metrics.log_metrics

tests/e2e/metrics/test_metrics.py

@@ -13,26 +13,30 @@ def config() -> conftest.LambdaConfig:
     return {
         "parameters": {},
         "environment_variables": {
-            "METRIC_NAMESPACE": f"powertools-e2e-metric-{uuid.uuid4()}",
-            "METRIC_NAME": "business-metric",
-            "SERVICE_NAME": "test-powertools-service",
+            "POWERTOOLS_METRICS_NAMESPACE": "powertools-e2e-metric",
+            "POWERTOOLS_SERVICE_NAME": f"test-powertools-service",
+            "METRIC_NAME": f"business-metric-{uuid.uuid4()}",
         },
     }
 
 
 @pytest.mark.e2e
 def test_basic_lambda_metric_visible(execute_lambda: conftest.InfrastructureOutput, config: conftest.LambdaConfig):
+    # GIVEN
     start_date = execute_lambda.get_lambda_execution_time()
     end_date = start_date + datetime.timedelta(minutes=5)
 
+    # WHEN
     metrics = helpers.get_metrics(
         start_date=start_date,
         end_date=end_date,
-        namespace=config["environment_variables"]["METRIC_NAMESPACE"],
+        namespace=config["environment_variables"]["POWERTOOLS_METRICS_NAMESPACE"],
         metric_name=config["environment_variables"]["METRIC_NAME"],
-        service_name=config["environment_variables"]["SERVICE_NAME"],
+        service_name=config["environment_variables"]["POWERTOOLS_SERVICE_NAME"],
         cw_client=boto3.client(service_name="cloudwatch"),
     )
+
+    # THEN
     assert metrics["Timestamps"] and len(metrics["Timestamps"]) == 1
     assert metrics["Values"] and len(metrics["Values"]) == 1
     assert metrics["Values"][0] == 1

tests/e2e/tracer/handlers/basic_handler.py

@@ -1,16 +1,25 @@
+import asyncio
 import os
 
 from aws_lambda_powertools import Tracer
+from aws_lambda_powertools.utilities.typing import LambdaContext
 
-tracer = Tracer()
 tracer = Tracer(service="e2e-tests-app")
 
 ANNOTATION_KEY = os.environ["ANNOTATION_KEY"]
 ANNOTATION_VALUE = os.environ["ANNOTATION_VALUE"]
+ANNOTATION_ASYNC_VALUE = os.environ["ANNOTATION_ASYNC_VALUE"]
 
 
 @tracer.capture_lambda_handler
-def lambda_handler(event, context):
-    tracer.put_annotation(key=ANNOTATION_KEY, value=ANNOTATION_VALUE)
+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

@@ -13,34 +13,43 @@
 def config():
     return {
         "parameters": {"tracing": "ACTIVE"},
-        "environment_variables": {"ANNOTATION_KEY": f"e2e-tracer-{uuid.uuid4()}", "ANNOTATION_VALUE": "stored"},
+        "environment_variables": {
+            "ANNOTATION_KEY": f"e2e-tracer-{uuid.uuid4()}",
+            "ANNOTATION_VALUE": "stored",
+            "ANNOTATION_ASYNC_VALUE": "payments",
+        },
     }
 
 
 @pytest.mark.e2e
-def test_basic_lambda_trace_visible(execute_lambda: conftest.InfrastructureOutput, config: conftest.LambdaConfig):
+def test_basic_lambda_async_trace_visible(execute_lambda: conftest.InfrastructureOutput, config: conftest.LambdaConfig):
+    # GIVEN
     lambda_arn = execute_lambda.get_lambda_arn(name="basichandlerarn")
     start_date = execute_lambda.get_lambda_execution_time()
     end_date = start_date + datetime.timedelta(minutes=5)
 
+    # WHEN
     trace = helpers.get_traces(
         start_date=start_date,
         end_date=end_date,
         lambda_function_name=lambda_arn.split(":")[-1],
         xray_client=boto3.client("xray"),
     )
 
-    for segment in trace["Traces"][0]["Segments"]:
-        document = json.loads(segment["Document"])
-        if document["origin"] == "AWS::Lambda::Function":
-            for subsegment in document["subsegments"]:
-                if subsegment["name"] == "Invocation":
-                    for x_subsegment in subsegment["subsegments"]:
-                        metadata = x_subsegment["metadata"]
-                        annotation = x_subsegment["annotations"]
+    # THEN
+    info = helpers.find_trace_additional_info(trace=trace)
+    print(info)
+    handler_trace_segment = [trace_segment for trace_segment in info if trace_segment.name == "## lambda_handler"][0]
+    collect_payment_trace_segment = [
+        trace_segment for trace_segment in info if trace_segment.name == "## collect_payment"
+    ][0]
 
+    assert handler_trace_segment.annotations["Service"] == "e2e-tests-app"
     assert (
-        metadata["e2e-tests-app"][config["environment_variables"]["ANNOTATION_KEY"]]
+        handler_trace_segment.metadata["e2e-tests-app"][config["environment_variables"]["ANNOTATION_KEY"]]
         == config["environment_variables"]["ANNOTATION_VALUE"]
     )
-    assert annotation["Service"] == "e2e-tests-app"
+    assert (
+        collect_payment_trace_segment.metadata["e2e-tests-app"][config["environment_variables"]["ANNOTATION_KEY"]]
+        == config["environment_variables"]["ANNOTATION_ASYNC_VALUE"]
+    )

tests/e2e/utils/helpers.py

@@ -1,7 +1,7 @@
 import json
 from datetime import datetime
 from functools import lru_cache
-from typing import Any, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 from pydantic import BaseModel
 from retry import retry
@@ -23,6 +23,12 @@ class Log(BaseModel):
     extra_info: Optional[str]
 
 
+class TraceSegment(BaseModel):
+    name: str
+    metadata: Dict = {}
+    annotations: Dict = {}
+
+
 def trigger_lambda(lambda_arn: str, client: Any):
     response = client.invoke(FunctionName=lambda_arn, InvocationType="RequestResponse")
     return response
@@ -76,7 +82,7 @@ def get_metrics(
 
 
 @retry(ValueError, delay=1, jitter=1, tries=10)
-def get_traces(lambda_function_name: str, xray_client: Any, start_date: datetime, end_date: datetime):
+def get_traces(lambda_function_name: str, xray_client: Any, start_date: datetime, end_date: datetime) -> Dict:
     paginator = xray_client.get_paginator("get_trace_summaries")
     response_iterator = paginator.paginate(
         StartTime=start_date,
@@ -95,3 +101,27 @@ def get_traces(lambda_function_name: str, xray_client: Any, start_date: datetime
     )
 
     return trace_details
+
+
+def find_trace_additional_info(trace: Dict) -> List[TraceSegment]:
+    info = []
+    for segment in trace["Traces"][0]["Segments"]:
+        document = json.loads(segment["Document"])
+        if document["origin"] == "AWS::Lambda::Function":
+            for subsegment in document["subsegments"]:
+                if subsegment["name"] == "Invocation":
+                    find_meta(segment=subsegment, result=info)
+    return info
+
+
+def find_meta(segment: dict, result: List):
+    for x_subsegment in segment["subsegments"]:
+        result.append(
+            TraceSegment(
+                name=x_subsegment["name"],
+                metadata=x_subsegment.get("metadata", {}),
+                annotations=x_subsegment.get("annotations", {}),
+            )
+        )
+        if x_subsegment.get("subsegments"):
+            find_meta(segment=x_subsegment, result=result)

tests/e2e/utils/infrastructure.py

@@ -55,11 +55,10 @@ def _create_layer(self, stack: Stack):
                     command=[
                         "bash",
                         "-c",
-                        f"export PYTHONPYCACHEPREFIX='/tmp/.cache/pycache/';\
-                            poetry export --with-credentials --format requirements.txt --output /tmp/requirements.txt &&\
+                        f"poetry export --with-credentials --format requirements.txt --output /tmp/requirements.txt &&\
                             pip install -r /tmp/requirements.txt -t {output_dir} &&\
                             cp -R {input_dir} {output_dir} &&\
-                            find {output_dir}/ -regex '^.*\\(__pycache__\\|\\.py[co]\\)$' -delete",
+                            find {output_dir}/ -regex '^.*__pycache__.*' -delete",
                     ],
                 ),
             ),
@@ -155,15 +154,11 @@ def _upload_assets(self, template: dict, asset_root_dir: str):
             buf.seek(0)
             self.s3_client.upload_fileobj(Fileobj=buf, Bucket=bucket, Key=s3_key)
 
-    def _find_files(self, directory: str, only_py: bool = False) -> list:
+    def _find_files(self, directory: str) -> list:
         file_paths = []
         for root, _, files in os.walk(directory):
             for filename in files:
-                if only_py:
-                    if filename.endswith(".py"):
-                        file_paths.append(os.path.join(root, filename))
-                else:
-                    file_paths.append(os.path.join(root, filename))
+                file_paths.append(os.path.join(root, filename))
         return file_paths
 
     def _deploy_stack(self, stack_name: str, template: dict):