Add tests

* Extract is_durable_context type guard
* Add tests to verify inject_lambda_context handles durable context
* Add tests to verify log_metrics handles durable context

Author: ConnorKirk

aws_lambda_powertools/logging/logger.py

@@ -38,6 +38,7 @@
 from aws_lambda_powertools.shared.functions import (
     extract_event_from_common_models,
     get_tracer_id,
+    is_durable_context,
     resolve_env_var_choice,
     resolve_truthy_env_var_choice,
 )
@@ -522,7 +523,7 @@ def handler(event, context):
         def decorate(event, context, *args, **kwargs):
             unwrapped_context = (
                 build_lambda_context_model(context.lambda_context)
-                if hasattr(context, "step")
+                if is_durable_context(context)
                 else build_lambda_context_model(context)
             )
 

aws_lambda_powertools/metrics/base.py

@@ -37,7 +37,7 @@
     reset_cold_start_flag,  # noqa: F401  # backwards compatibility
 )
 from aws_lambda_powertools.shared import constants
-from aws_lambda_powertools.shared.functions import resolve_env_var_choice
+from aws_lambda_powertools.shared.functions import is_durable_context, resolve_env_var_choice
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Generator
@@ -430,7 +430,7 @@ def handler(event, context):
 
         @functools.wraps(lambda_handler)
         def decorate(event, context, *args, **kwargs):
-            unwrapped_context = context.lambda_context if hasattr(context, "step") else context
+            unwrapped_context = context.lambda_context if is_durable_context(context) else context
             try:
                 if default_dimensions:
                     self.set_default_dimensions(**default_dimensions)

aws_lambda_powertools/metrics/provider/base.py

@@ -6,6 +6,7 @@
 from typing import TYPE_CHECKING, Any
 
 from aws_lambda_powertools.metrics.provider import cold_start
+from aws_lambda_powertools.shared.functions import is_durable_context
 
 if TYPE_CHECKING:
     from aws_lambda_powertools.shared.types import AnyCallableT
@@ -206,7 +207,7 @@ def decorate(event, context, *args, **kwargs):
             try:
                 response = lambda_handler(event, context, *args, **kwargs)
                 if capture_cold_start_metric:
-                    unwrapped_context = context.lambda_context if hasattr(context, "step") else context
+                    unwrapped_context = context.lambda_context if is_durable_context(context) else context
                     self._add_cold_start_metric(context=unwrapped_context)
             finally:
                 self.flush_metrics(raise_on_empty_metrics=raise_on_empty_metrics)

aws_lambda_powertools/shared/functions.py

@@ -8,13 +8,15 @@
 import warnings
 from binascii import Error as BinAsciiError
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, overload
+from typing import TYPE_CHECKING, Any, TypeGuard, overload
 
 from aws_lambda_powertools.shared import constants
 
 if TYPE_CHECKING:
     from collections.abc import Generator
 
+    from aws_lambda_powertools.utilities.typing import DurableContextProtocol
+
 logger = logging.getLogger(__name__)
 
 
@@ -307,3 +309,8 @@ def decode_header_bytes(byte_list):
     # Convert signed bytes to unsigned (0-255 range)
     unsigned_bytes = [(b & 0xFF) for b in byte_list]
     return bytes(unsigned_bytes)
+
+
+def is_durable_context(context: Any) -> TypeGuard[DurableContextProtocol]:
+    """Check if context is a Step Functions durable context wrapping a Lambda context."""
+    return hasattr(context, "state") and hasattr(context, "lambda_context")

tests/functional/logger/required_dependencies/test_logger.py

@@ -48,6 +48,11 @@ def lambda_context():
     return namedtuple("LambdaContext", lambda_context.keys())(*lambda_context.values())
 
 
+@pytest.fixture
+def durable_context(lambda_context):
+    return namedtuple("DurableContext", ["state", "lambda_context"])(state={}, lambda_context=lambda_context)
+
+
 @pytest.fixture
 def lambda_event():
     return {"greeting": "hello"}
@@ -1578,3 +1583,20 @@ def test_child_logger_with_caplog(caplog):
 
     assert len(caplog.records) == 1
     assert pytest_handler_existence is True
+
+
+def test_logger_with_durable_context(lambda_context, durable_context, stdout, service_name):
+    # GIVEN Logger is initialized and a durable context wrapping the lambda context
+    logger = Logger(service=service_name, stream=stdout)
+
+    @logger.inject_lambda_context
+    def handler(event, context):
+        logger.info("Hello")
+
+    # WHEN handler is called with durable context
+    handler({}, durable_context)
+
+    # THEN lambda contextual info should be extracted from durable context
+    log = capture_logging_output(stdout)
+    assert log["function_name"] == lambda_context.function_name
+    assert log["function_request_id"] == lambda_context.aws_request_id

tests/functional/metrics/conftest.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections import namedtuple
 from typing import Any
 
 import pytest
@@ -96,3 +97,13 @@ def a_hundred_metrics() -> list[dict[str, str]]:
 @pytest.fixture
 def a_hundred_metric_values() -> list[dict[str, str]]:
     return [{"name": "metric", "unit": "Count", "value": i} for i in range(100)]
+
+
+@pytest.fixture
+def lambda_context():
+    return namedtuple("LambdaContext", "function_name")("example_fn")
+
+
+@pytest.fixture
+def durable_context(lambda_context):
+    return namedtuple("DurableContext", ["state", "lambda_context"])(state={}, lambda_context=lambda_context)

tests/functional/metrics/required_dependencies/test_metrics_cloudwatch_emf.py

@@ -1573,3 +1573,42 @@ def test_metrics_disabled_with_dev_mode_false_and_metrics_disabled_true(monkeypa
     # THEN no metrics should have been recorded
     captured = capsys.readouterr()
     assert not captured.out
+
+
+def test_log_metrics_with_durable_context(capsys, metrics, dimensions, namespace, durable_context):
+    # GIVEN Metrics is initialized and a durable context wrapping the lambda context
+    my_metrics = Metrics(namespace=namespace)
+    for metric in metrics:
+        my_metrics.add_metric(**metric)
+    for dimension in dimensions:
+        my_metrics.add_dimension(**dimension)
+
+    @my_metrics.log_metrics
+    def lambda_handler(evt, ctx):
+        pass
+
+    # WHEN handler is called with durable context
+    lambda_handler({}, durable_context)
+    output = capture_metrics_output(capsys)
+    expected = serialize_metrics(metrics=metrics, dimensions=dimensions, namespace=namespace)
+
+    # THEN metrics should be flushed correctly
+    remove_timestamp(metrics=[output, expected])
+    assert expected == output
+
+
+def test_log_metrics_capture_cold_start_metric_with_durable_context(capsys, namespace, service, durable_context):
+    # GIVEN Metrics is initialized and a durable context wrapping the lambda context
+    my_metrics = Metrics(service=service, namespace=namespace)
+
+    @my_metrics.log_metrics(capture_cold_start_metric=True)
+    def lambda_handler(evt, context):
+        pass
+
+    # WHEN handler is called with durable context
+    lambda_handler({}, durable_context)
+    output = capture_metrics_output(capsys)
+
+    # THEN ColdStart metric should use function_name from unwrapped lambda context
+    assert output["ColdStart"] == [1.0]
+    assert output["function_name"] == "example_fn"