From 879d627f6908afe09919cdfc95cafd506a58bf16 Mon Sep 17 00:00:00 2001 From: heitorlessa Date: Mon, 29 Nov 2021 11:58:00 +0100 Subject: [PATCH] feat(tracer): add ColdStart annotation as false when it's a warm start --- aws_lambda_powertools/tracing/tracer.py | 5 +++-- tests/unit/test_tracing.py | 29 ++++++++++++++++++++----- 2 files changed, 27 insertions(+), 7 deletions(-) diff --git a/aws_lambda_powertools/tracing/tracer.py b/aws_lambda_powertools/tracing/tracer.py index 2626793304c..2beab0483be 100644 --- a/aws_lambda_powertools/tracing/tracer.py +++ b/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: diff --git a/tests/unit/test_tracing.py b/tests/unit/test_tracing.py index 2b147ec4405..7c8b6244f01 100644 --- a/tests/unit/test_tracing.py +++ b/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)