fix(tracer): add warm start annotation (ColdStart=False) (aws-powertools#851)

Author: heitorlessa

aws_lambda_powertools/tracing/tracer.py

@@ -304,9 +304,10 @@ def handler(event, context):
  def decorate(event, context, **kwargs):
  with self.provider.in_subsegment(name=f"## {lambda_handler_name}") as subsegment:
  global is_cold_start
+ logger.debug("Annotating cold start")
+ subsegment.put_annotation(key="ColdStart", value=is_cold_start)
+
  if is_cold_start:
- logger.debug("Annotating cold start")
- subsegment.put_annotation(key="ColdStart", value=True)
  is_cold_start = False
 
  try:

tests/unit/test_tracing.py

@@ -51,7 +51,9 @@ def patch(self, *args, **kwargs):
 def reset_tracing_config(mocker):
  Tracer._reset_config()
  # reset global cold start module
- mocker.patch("aws_lambda_powertools.tracing.tracer.is_cold_start", return_value=True)
+ mocker.patch(
+ "aws_lambda_powertools.tracing.tracer.is_cold_start", new_callable=mocker.PropertyMock(return_value=True)
+ )
  yield
 
 
@@ -79,7 +81,7 @@ class InSubsegment(NamedTuple):
  yield in_subsegment
 
 
-def test_tracer_lambda_handler(mocker, dummy_response, provider_stub, in_subsegment_mock):
+def test_tracer_lambda_handler_subsegment(mocker, dummy_response, provider_stub, in_subsegment_mock):
  # GIVEN Tracer is initialized with booking as the service name
  provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
  tracer = Tracer(provider=provider, service="booking")
@@ -92,15 +94,13 @@ def handler(event, context):
  handler({}, mocker.MagicMock())
 
  # THEN we should have a subsegment named handler
- # annotate cold start, and add its response as trace metadata
+ # add its response as trace metadata
  # and use service name as a metadata namespace
  assert in_subsegment_mock.in_subsegment.call_count == 1
  assert in_subsegment_mock.in_subsegment.call_args == mocker.call(name="## handler")
  assert in_subsegment_mock.put_metadata.call_args == mocker.call(
  key="handler response", value=dummy_response, namespace="booking"
  )
- assert in_subsegment_mock.put_annotation.call_count == 1
- assert in_subsegment_mock.put_annotation.call_args == mocker.call(key="ColdStart", value=True)
 
 
 def test_tracer_method(mocker, dummy_response, provider_stub, in_subsegment_mock):
@@ -571,3 +571,22 @@ def greeting(name, message):
 
  # THEN we should not add any metadata
  assert in_subsegment_mock.put_metadata.call_count == 0
+
+
+def test_tracer_lambda_handler_cold_start(mocker, dummy_response, provider_stub, in_subsegment_mock):
+ # GIVEN
+ provider = provider_stub(in_subsegment=in_subsegment_mock.in_subsegment)
+ tracer = Tracer(provider=provider, service="booking")
+
+ # WHEN
+ @tracer.capture_lambda_handler
+ def handler(event, context):
+ return dummy_response
+
+ handler({}, mocker.MagicMock())
+
+ # THEN
+ assert in_subsegment_mock.put_annotation.call_args == mocker.call(key="ColdStart", value=True)
+
+ handler({}, mocker.MagicMock())
+ assert in_subsegment_mock.put_annotation.call_args == mocker.call(key="ColdStart", value=False)