use initialization type to determine readiness (#174)

Author: saartochner-lumigo

src/lumigo_tracer/lumigo_utils.py

@@ -661,3 +661,7 @@ def is_error_code(status_code: int) -> bool:
 
 def is_aws_arn(string_to_validate: Optional[str]) -> bool:
     return bool(string_to_validate and string_to_validate.startswith("arn:aws:"))
+
+
+def is_provision_concurrency_initialization() -> bool:
+    return os.environ.get("AWS_LAMBDA_INITIALIZATION_TYPE") == "provisioned-concurrency"

src/lumigo_tracer/spans_container.py

@@ -23,6 +23,7 @@
     create_step_function_span,
     get_current_ms_time,
     get_region,
+    is_provision_concurrency_initialization,
 )
 from lumigo_tracer import lumigo_utils
 from lumigo_tracer.parsing_utils import parse_trace_id, safe_split_get, recursive_json_join
@@ -78,6 +79,7 @@ def __init__(
             "info": {"tracer": {"version": version}, "traceId": {"Root": trace_root}},
             "token": Configuration.token,
         }
+        is_cold = SpansContainer.is_cold and not is_provision_concurrency_initialization()
         self.function_span = recursive_json_join(
             {
                 "id": request_id,
@@ -87,7 +89,7 @@ def __init__(
                 "event": event,
                 "envs": envs,
                 "memoryAllocated": memory_allocated,
-                "readiness": "cold" if SpansContainer.is_cold else "warm",
+                "readiness": "cold" if is_cold else "warm",
                 "info": {
                     "logStreamName": log_stream_name,
                     "logGroupName": log_group_name,

src/test/unit/test_tracer.py

@@ -96,6 +96,19 @@ def lambda_test_function(event, context):
     assert span["return_value"] == '{"a": [1.0]}'
 
 
+def test_lambda_wrapper_provision_concurrency_is_warm(context, monkeypatch):
+    monkeypatch.setattr(SpansContainer, "is_cold", True)
+    monkeypatch.setenv("AWS_LAMBDA_INITIALIZATION_TYPE", "provisioned-concurrency")
+
+    @lumigo_tracer(token="123")
+    def lambda_test_function(event, context):
+        return {"a": "b"}
+
+    lambda_test_function({}, context)
+    span = SpansContainer.get_span().function_span
+    assert span["readiness"] == "warm"
+
+
 def test_kill_switch(monkeypatch, context):
     monkeypatch.setattr(os, "environ", {"LUMIGO_SWITCH_OFF": "true"})