chore(tests): refactor E2E tracer to ease maintenance, writing tests and parallelization

.flake8

@@ -3,6 +3,9 @@ exclude = docs, .eggs, setup.py, example, .aws-sam, .git, dist, *.md, *.yaml, ex
 ignore = E203, E266, W503, BLK100, W291, I004
 max-line-length = 120
 max-complexity = 15
+per-file-ignores =
+    tests/e2e/utils/data_builder/__init__.py:F401
+    tests/e2e/utils/data_fetcher/__init__.py:F401
 
 [isort]
 multi_line_output = 3

tests/e2e/conftest.py

@@ -5,6 +5,8 @@
 
 import boto3
 
+from tests.e2e.utils import data_fetcher, infrastructure
+
 # We only need typing_extensions for python versions <3.8
 if sys.version_info >= (3, 8):
     from typing import TypedDict
@@ -14,7 +16,6 @@
 from typing import Dict, Generator, Optional
 
 import pytest
-from e2e.utils import helpers, infrastructure
 
 
 class LambdaConfig(TypedDict):
@@ -61,5 +62,5 @@ def execute_lambda(create_infrastructure) -> InfrastructureOutput:
     session = boto3.Session()
     client = session.client("lambda")
     for _, arn in create_infrastructure.items():
-        helpers.trigger_lambda(lambda_arn=arn, client=client)
+        data_fetcher.get_lambda_response(lambda_arn=arn, client=client)
     return InfrastructureOutput(arns=create_infrastructure, execution_time=execution_time)

tests/e2e/logger/test_logger.py

@@ -1,7 +1,8 @@
 import boto3
 import pytest
 from e2e import conftest
-from e2e.utils import helpers
+
+from tests.e2e.utils import data_fetcher
 
 
 @pytest.fixture(scope="module")
@@ -23,7 +24,7 @@ def test_basic_lambda_logs_visible(execute_lambda: conftest.InfrastructureOutput
     cw_client = boto3.client("logs")
 
     # WHEN
-    filtered_logs = helpers.get_logs(lambda_function_name=lambda_name, start_time=timestamp, log_client=cw_client)
+    filtered_logs = data_fetcher.get_logs(lambda_function_name=lambda_name, start_time=timestamp, log_client=cw_client)
 
     # THEN
     assert any(
@@ -42,7 +43,7 @@ def test_basic_lambda_no_debug_logs_visible(
     cw_client = boto3.client("logs")
 
     # WHEN
-    filtered_logs = helpers.get_logs(lambda_function_name=lambda_name, start_time=timestamp, log_client=cw_client)
+    filtered_logs = data_fetcher.get_logs(lambda_function_name=lambda_name, start_time=timestamp, log_client=cw_client)
 
     # THEN
     assert not any(
@@ -66,7 +67,7 @@ def test_basic_lambda_contextual_data_logged(execute_lambda: conftest.Infrastruc
     cw_client = boto3.client("logs")
 
     # WHEN
-    filtered_logs = helpers.get_logs(lambda_function_name=lambda_name, start_time=timestamp, log_client=cw_client)
+    filtered_logs = data_fetcher.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)
@@ -81,7 +82,7 @@ def test_basic_lambda_additional_key_persistence_basic_lambda(
     cw_client = boto3.client("logs")
 
     # WHEN
-    filtered_logs = helpers.get_logs(lambda_function_name=lambda_name, start_time=timestamp, log_client=cw_client)
+    filtered_logs = data_fetcher.get_logs(lambda_function_name=lambda_name, start_time=timestamp, log_client=cw_client)
 
     # THEN
     assert any(
@@ -100,7 +101,7 @@ def test_basic_lambda_empty_event_logged(execute_lambda: conftest.Infrastructure
     cw_client = boto3.client("logs")
 
     # WHEN
-    filtered_logs = helpers.get_logs(lambda_function_name=lambda_name, start_time=timestamp, log_client=cw_client)
+    filtered_logs = data_fetcher.get_logs(lambda_function_name=lambda_name, start_time=timestamp, log_client=cw_client)
 
     # THEN
     assert any(log.message == {} for log in filtered_logs)
@@ -122,7 +123,7 @@ def test_no_context_lambda_contextual_data_not_logged(execute_lambda: conftest.I
     cw_client = boto3.client("logs")
 
     # WHEN
-    filtered_logs = helpers.get_logs(lambda_function_name=lambda_name, start_time=timestamp, log_client=cw_client)
+    filtered_logs = data_fetcher.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)
@@ -136,7 +137,7 @@ def test_no_context_lambda_event_not_logged(execute_lambda: conftest.Infrastruct
     cw_client = boto3.client("logs")
 
     # WHEN
-    filtered_logs = helpers.get_logs(lambda_function_name=lambda_name, start_time=timestamp, log_client=cw_client)
+    filtered_logs = data_fetcher.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/conftest.py

@@ -10,8 +10,12 @@ def infrastructure(request: pytest.FixtureRequest, tmp_path_factory: pytest.Temp
 
     Parameters
     ----------
-    request : fixtures.SubRequest
-        test fixture containing metadata about test execution
+    request : pytest.FixtureRequest
+        pytest request fixture to introspect absolute path to test being executed
+    tmp_path_factory : pytest.TempPathFactory
+        pytest temporary path factory to discover shared tmp when multiple CPU processes are spun up
+    worker_id : str
+        pytest-xdist worker identification to detect whether parallelization is enabled
 
     Yields
     ------

tests/e2e/metrics/test_metrics.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from tests.e2e.utils import helpers
+from tests.e2e.utils import data_builder, data_fetcher
 
 
 @pytest.fixture
@@ -30,40 +30,40 @@ def cold_start_fn_arn(infrastructure: dict) -> str:
 
 def test_basic_lambda_metric_is_visible(basic_handler_fn: str, basic_handler_fn_arn: str):
     # GIVEN
-    metric_name = helpers.build_metric_name()
-    service = helpers.build_service_name()
-    dimensions = helpers.build_add_dimensions_input(service=service)
-    metrics = helpers.build_multiple_add_metric_input(metric_name=metric_name, value=1, quantity=3)
+    metric_name = data_builder.build_metric_name()
+    service = data_builder.build_service_name()
+    dimensions = data_builder.build_add_dimensions_input(service=service)
+    metrics = data_builder.build_multiple_add_metric_input(metric_name=metric_name, value=1, quantity=3)
 
     # WHEN
     event = json.dumps({"metrics": metrics, "service": service, "namespace": METRIC_NAMESPACE})
-    _, execution_time = helpers.trigger_lambda(lambda_arn=basic_handler_fn_arn, payload=event)
+    _, execution_time = data_fetcher.get_lambda_response(lambda_arn=basic_handler_fn_arn, payload=event)
 
-    metrics = helpers.get_metrics(
+    my_metrics = data_fetcher.get_metrics(
         namespace=METRIC_NAMESPACE, start_date=execution_time, metric_name=metric_name, dimensions=dimensions
     )
 
     # THEN
-    metric_data = metrics.get("Values", [])
+    metric_data = my_metrics.get("Values", [])
     assert metric_data and metric_data[0] == 3.0
 
 
 def test_cold_start_metric(cold_start_fn_arn: str, cold_start_fn: str):
     # GIVEN
     metric_name = "ColdStart"
-    service = helpers.build_service_name()
-    dimensions = helpers.build_add_dimensions_input(function_name=cold_start_fn, service=service)
+    service = data_builder.build_service_name()
+    dimensions = data_builder.build_add_dimensions_input(function_name=cold_start_fn, service=service)
 
     # WHEN we invoke twice
     event = json.dumps({"service": service, "namespace": METRIC_NAMESPACE})
 
-    _, execution_time = helpers.trigger_lambda(lambda_arn=cold_start_fn_arn, payload=event)
-    _, _ = helpers.trigger_lambda(lambda_arn=cold_start_fn_arn, payload=event)
+    _, execution_time = data_fetcher.get_lambda_response(lambda_arn=cold_start_fn_arn, payload=event)
+    _, _ = data_fetcher.get_lambda_response(lambda_arn=cold_start_fn_arn, payload=event)
 
-    metrics = helpers.get_metrics(
+    my_metrics = data_fetcher.get_metrics(
         namespace=METRIC_NAMESPACE, start_date=execution_time, metric_name=metric_name, dimensions=dimensions
     )
 
     # THEN
-    metric_data = metrics.get("Values", [])
+    metric_data = my_metrics.get("Values", [])
     assert metric_data and metric_data[0] == 1.0

tests/e2e/tracer/conftest.py

@@ -0,0 +1,25 @@
+import pytest
+
+from tests.e2e.tracer.infrastructure import TracerStack
+from tests.e2e.utils.infrastructure import deploy_once
+
+
+@pytest.fixture(autouse=True, scope="module")
+def infrastructure(request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory, worker_id: str):
+    """Setup and teardown logic for E2E test infrastructure
+
+    Parameters
+    ----------
+    request : pytest.FixtureRequest
+        pytest request fixture to introspect absolute path to test being executed
+    tmp_path_factory : pytest.TempPathFactory
+        pytest temporary path factory to discover shared tmp when multiple CPU processes are spun up
+    worker_id : str
+        pytest-xdist worker identification to detect whether parallelization is enabled
+
+    Yields
+    ------
+    Dict[str, str]
+        CloudFormation Outputs from deployed infrastructure
+    """
+    yield from deploy_once(stack=TracerStack, request=request, tmp_path_factory=tmp_path_factory, worker_id=worker_id)

tests/e2e/tracer/handlers/async_capture.py

@@ -0,0 +1,16 @@
+import asyncio
+from uuid import uuid4
+
+from aws_lambda_powertools import Tracer
+from aws_lambda_powertools.utilities.typing import LambdaContext
+
+tracer = Tracer()
+
+
+@tracer.capture_method
+async def async_get_users():
+    return [{"id": f"{uuid4()}"} for _ in range(5)]
+
+
+def lambda_handler(event: dict, context: LambdaContext):
+    return asyncio.run(async_get_users())

tests/e2e/tracer/handlers/basic_handler.py

@@ -1,25 +1,16 @@
-import asyncio
-import os
+from uuid import uuid4
 
 from aws_lambda_powertools import Tracer
 from aws_lambda_powertools.utilities.typing import LambdaContext
 
-tracer = Tracer(service="e2e-tests-app")
+tracer = Tracer()
 
-ANNOTATION_KEY = os.environ["ANNOTATION_KEY"]
-ANNOTATION_VALUE = os.environ["ANNOTATION_VALUE"]
-ANNOTATION_ASYNC_VALUE = os.environ["ANNOTATION_ASYNC_VALUE"]
+
+@tracer.capture_method
+def get_todos():
+    return [{"id": f"{uuid4()}", "completed": False} for _ in range(5)]
 
 
 @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"
+    return get_todos()

tests/e2e/tracer/infrastructure.py

@@ -0,0 +1,17 @@
+from pathlib import Path
+
+from tests.e2e.utils.data_builder import build_service_name
+from tests.e2e.utils.infrastructure import BaseInfrastructureV2
+
+
+class TracerStack(BaseInfrastructureV2):
+    # Maintenance: Tracer doesn't support dynamic service injection (tracer.py L310)
+    # we could move after handler response or adopt env vars usage in e2e tests
+    SERVICE_NAME: str = build_service_name()
+
+    def __init__(self, handlers_dir: Path, feature_name: str = "tracer") -> None:
+        super().__init__(feature_name, handlers_dir)
+
+    def create_resources(self) -> None:
+        env_vars = {"POWERTOOLS_SERVICE_NAME": self.SERVICE_NAME}
+        self.create_lambda_functions(function_props={"environment": env_vars})

tests/e2e/tracer/test_tracer.py

@@ -1,51 +1,69 @@
-import datetime
-import uuid
-
-import boto3
 import pytest
-from e2e import conftest
-from e2e.utils import helpers
+
+from tests.e2e.tracer.handlers import async_capture, basic_handler
+from tests.e2e.tracer.infrastructure import TracerStack
+from tests.e2e.utils import data_builder, data_fetcher
+
+
+@pytest.fixture
+def basic_handler_fn_arn(infrastructure: dict) -> str:
+    return infrastructure.get("BasicHandlerArn", "")
+
+
+@pytest.fixture
+def basic_handler_fn(infrastructure: dict) -> str:
+    return infrastructure.get("BasicHandler", "")
+
+
+@pytest.fixture
+def async_fn_arn(infrastructure: dict) -> str:
+    return infrastructure.get("AsyncCaptureArn", "")
+
+
+@pytest.fixture
+def async_fn(infrastructure: dict) -> str:
+    return infrastructure.get("AsyncCapture", "")
 
 
-@pytest.fixture(scope="module")
-def config() -> conftest.LambdaConfig:
-    return {
-        "parameters": {"tracing": "ACTIVE"},
-        "environment_variables": {
-            "ANNOTATION_KEY": f"e2e-tracer-{str(uuid.uuid4()).replace('-','_')}",
-            "ANNOTATION_VALUE": "stored",
-            "ANNOTATION_ASYNC_VALUE": "payments",
-        },
-    }
+def test_lambda_handler_trace_is_visible(basic_handler_fn_arn: str, basic_handler_fn: str):
+    # GIVEN
+    handler_name = basic_handler.lambda_handler.__name__
+    handler_subsegment = f"## {handler_name}"
+    handler_metadata_key = f"{handler_name} response"
+
+    method_name = basic_handler.get_todos.__name__
+    method_subsegment = f"## {method_name}"
+    handler_metadata_key = f"{method_name} response"
+
+    trace_query = data_builder.build_trace_default_query(function_name=basic_handler_fn)
+
+    # WHEN
+    _, execution_time = data_fetcher.get_lambda_response(lambda_arn=basic_handler_fn_arn)
+    data_fetcher.get_lambda_response(lambda_arn=basic_handler_fn_arn)
+
+    # THEN
+    trace = data_fetcher.get_traces(start_date=execution_time, filter_expression=trace_query, minimum_traces=2)
+
+    assert len(trace.get_annotation(key="ColdStart", value=True)) == 1
+    assert len(trace.get_metadata(key=handler_metadata_key, namespace=TracerStack.SERVICE_NAME)) == 2
+    assert len(trace.get_metadata(key=handler_metadata_key, namespace=TracerStack.SERVICE_NAME)) == 2
+    assert len(trace.get_subsegment(name=handler_subsegment)) == 2
+    assert len(trace.get_subsegment(name=method_subsegment)) == 2
 
 
-def test_basic_lambda_async_trace_visible(execute_lambda: conftest.InfrastructureOutput, config: conftest.LambdaConfig):
+def test_async_trace_is_visible(async_fn_arn: str, async_fn: str):
     # GIVEN
-    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_name}")'
+    async_fn_name = async_capture.async_get_users.__name__
+    async_fn_name_subsegment = f"## {async_fn_name}"
+    async_fn_name_metadata_key = f"{async_fn_name} response"
+
+    trace_query = data_builder.build_trace_default_query(function_name=async_fn)
 
     # WHEN
-    trace = helpers.get_traces(
-        start_date=start_date,
-        end_date=end_date,
-        filter_expression=trace_filter_exporession,
-        xray_client=boto3.client("xray"),
-    )
+    _, execution_time = data_fetcher.get_lambda_response(lambda_arn=async_fn_arn)
 
     # 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]
-
-    annotation_key = config["environment_variables"]["ANNOTATION_KEY"]
-    expected_value = config["environment_variables"]["ANNOTATION_VALUE"]
-    expected_async_value = config["environment_variables"]["ANNOTATION_ASYNC_VALUE"]
-
-    assert handler_trace_segment.annotations["Service"] == "e2e-tests-app"
-    assert handler_trace_segment.metadata["e2e-tests-app"][annotation_key] == expected_value
-    assert collect_payment_trace_segment.metadata["e2e-tests-app"][annotation_key] == expected_async_value
+    trace = data_fetcher.get_traces(start_date=execution_time, filter_expression=trace_query)
+
+    assert len(trace.get_subsegment(name=async_fn_name_subsegment)) == 1
+    assert len(trace.get_metadata(key=async_fn_name_metadata_key, namespace=TracerStack.SERVICE_NAME)) == 1