Rd 9297 disable addexecutiontag apis when lambda function not traced (#264)

* add execution_tag only if lambda is traced

* add execution_tag only if lambda is traced

* add execution_tag only if lambda is traced

* fix existing unit tests

Co-authored-by: CircleCI <no-reply@build.com>

Author: vladlLumigo

src/lumigo_tracer/lumigo_utils.py

@@ -860,3 +860,7 @@ def is_python_37() -> bool:
     InternalState.mark_timeout_to_edge()
 except Exception:
     pass
+
+
+def is_lambda_traced() -> bool:
+    return (not is_kill_switch_on()) and is_aws_environment()

src/lumigo_tracer/user_utils.py

@@ -4,7 +4,7 @@
 from typing import Dict, Optional
 
 from lumigo_tracer.spans_container import SpansContainer
-from lumigo_tracer.lumigo_utils import warn_client
+from lumigo_tracer.lumigo_utils import warn_client, is_lambda_traced
 
 LUMIGO_REPORT_ERROR_STRING = "[LUMIGO_LOG]"
 MAX_TAGS = 50
@@ -139,12 +139,17 @@ def add_execution_tag(key: str, value: str, should_log_errors: bool = True) -> b
     :param should_log_errors: Should a log message be printed in case the tag can't be added.
     """
     try:
-        key = str(key)
-        value = str(value)
-        tags_len = SpansContainer.get_span().get_tags_len()
-        if validate_tag(key, value, tags_len, should_log_errors):
-            SpansContainer.get_span().add_tag(key, value)
+        if is_lambda_traced():
+            key = str(key)
+            value = str(value)
+            tags_len = SpansContainer.get_span().get_tags_len()
+            if validate_tag(key, value, tags_len, should_log_errors):
+                SpansContainer.get_span().add_tag(key, value)
+            else:
+                return False
         else:
+            if should_log_errors:
+                warn_client(f"{ADD_TAG_ERROR_MSG_PREFIX}: lambda is not traced")
             return False
     except Exception:
         if should_log_errors:

src/test/conftest.py

@@ -121,3 +121,9 @@ def aws_env(monkeypatch):
         "_X_AMZN_TRACE_ID",
         "Root=1-12345678-111111111111111111111111;Parent=blablablablabla;Sampled=0",
     )
+
+
+@pytest.fixture
+def lambda_traced(monkeypatch):
+    monkeypatch.setenv("AWS_LAMBDA_FUNCTION_VERSION", "true")
+    monkeypatch.setenv("LUMIGO_SWITCH_OFF", "false")

src/test/unit/auto_tag/test_auto_tag_event.py

@@ -159,7 +159,7 @@ def test_auto_tag_key_not_in_header(monkeypatch):
     assert SpansContainer.get_span().execution_tags == []
 
 
-def test_auto_tag_key_in_header(monkeypatch):
+def test_auto_tag_key_in_header(monkeypatch, lambda_traced):
     set_header_key(monkeypatch, "Accept")
 
     event = {
@@ -323,7 +323,7 @@ def test_configuration_handler_is_supported(config, event, expected):
         ),
     ],
 )
-def test_configuration_handler_auto_tag(auto_tag_keys, event, result_tags):
+def test_configuration_handler_auto_tag(auto_tag_keys, event, result_tags, lambda_traced):
     Configuration.auto_tag = auto_tag_keys
     ConfigurationHandler.auto_tag(event)
     tags = SpansContainer.get_span().execution_tags
@@ -332,7 +332,7 @@ def test_configuration_handler_auto_tag(auto_tag_keys, event, result_tags):
         assert tag in tags
 
 
-def test_configuration_handler_auto_tag_failure(capsys):
+def test_configuration_handler_auto_tag_failure(capsys, lambda_traced):
     Configuration.auto_tag = [None, "key2"]
     ConfigurationHandler.auto_tag({"key1": datetime, "key2": "value"})
     tags = SpansContainer.get_span().execution_tags
@@ -347,7 +347,7 @@ def test_configuration_handler_auto_tag_failure(capsys):
         ({"a": "b"}, "{'a': 'b'}"),  # dict
     ],
 )
-def test_configuration_handler_auto_tag_non_string(value, expected):
+def test_configuration_handler_auto_tag_non_string(value, expected, lambda_traced):
     Configuration.auto_tag = ["key1"]
 
     ConfigurationHandler.auto_tag({"key1": value})

src/test/unit/test_spans_container.py

@@ -202,7 +202,7 @@ def test_timeout_mechanism_timeout_occurred_send_new_spans(monkeypatch, context,
 
 
 def test_timeout_mechanism_timeout_occurred_but_finish_check_enrichment(
-    monkeypatch, context, dummy_span, reporter_mock
+    monkeypatch, context, dummy_span, reporter_mock, lambda_traced
 ):
     SpansContainer.create_span()
     SpansContainer.get_span().start(context=context)

src/test/unit/test_tracer.py

@@ -423,7 +423,7 @@ def get_enrichment_spans(reporter_mock):
     return [s for s in final_send if s["type"] == ENRICHMENT_TYPE]
 
 
-def test_wrapping_with_tags(context, reporter_mock):
+def test_wrapping_with_tags(context, reporter_mock, lambda_traced):
     key = "my_key"
     value = "my_value"
 
@@ -443,7 +443,7 @@ def lambda_test_function(event, context):
     "key, event",
     [("my_key", {"my_key": "my_value"}), ("my_key.key2", {"my_key": {"key2": "my_value"}})],
 )
-def test_wrapping_with_auto_tags(context, key, event, reporter_mock):
+def test_wrapping_with_auto_tags(context, key, event, reporter_mock, lambda_traced):
     @lumigo_tracer(auto_tag=[key])
     def lambda_test_function(event, context):
         return "ret_value"

src/test/unit/test_user_utils.py

@@ -1,4 +1,5 @@
 import time
+import pytest
 
 from lumigo_tracer.spans_container import SpansContainer
 from lumigo_tracer.user_utils import (
@@ -103,7 +104,7 @@ def test_basic_info_warn_error(capsys):
     assert captured[2] == error_msg
 
 
-def test_add_execution_tag():
+def test_add_execution_tag(lambda_traced):
     key = "my_key"
     value = "my_value"
     assert add_execution_tag(key, value) is True
@@ -116,31 +117,31 @@ def test_start_manual_trace_simple_flow():
     assert SpansContainer.get_span().function_span[MANUAL_TRACES_KEY]
 
 
-def test_add_execution_key_tag_empty(capsys):
+def test_add_execution_key_tag_empty(capsys, lambda_traced):
     assert add_execution_tag("", "value") is False
     assert "Unable to add tag: key length" in capsys.readouterr().out
     assert SpansContainer.get_span().execution_tags == []
 
 
-def test_add_execution_value_tag_empty(capsys):
+def test_add_execution_value_tag_empty(capsys, lambda_traced):
     assert add_execution_tag("key", "") is False
     assert "Unable to add tag: value length" in capsys.readouterr().out
     assert SpansContainer.get_span().execution_tags == []
 
 
-def test_add_execution_tag_key_pass_max_chars(capsys):
+def test_add_execution_tag_key_pass_max_chars(capsys, lambda_traced):
     assert add_execution_tag("k" * (MAX_TAG_KEY_LEN + 1), "value") is False
     assert "Unable to add tag: key length" in capsys.readouterr().out
     assert SpansContainer.get_span().execution_tags == []
 
 
-def test_add_execution_tag_value_pass_max_chars(capsys):
+def test_add_execution_tag_value_pass_max_chars(capsys, lambda_traced):
     assert add_execution_tag("key", "v" * (MAX_TAG_VALUE_LEN + 1)) is False
     assert "Unable to add tag: value length" in capsys.readouterr().out
     assert SpansContainer.get_span().execution_tags == []
 
 
-def test_add_execution_tag_pass_max_tags(capsys):
+def test_add_execution_tag_pass_max_tags(capsys, lambda_traced):
     key = "my_key"
     value = "my_value"
 
@@ -166,3 +167,34 @@ def __str__(self):
     assert add_execution_tag("key", ExceptionOnStr()) is False
     assert "Unable to add tag" in capsys.readouterr().out
     assert SpansContainer.get_span().execution_tags == []
+
+
+@pytest.mark.parametrize(
+    ["kill_switch_value", "is_aws_environment_value", "expected_ret_value", "expected_tags"],
+    [
+        (  # happy flow - lambda is traced
+            "false",
+            "true",
+            True,
+            [{"key": "key", "value": "my-value"}],
+        ),
+        ("true", "true", False, []),  # kill switch on, is_aws_env true
+        ("true", "", False, []),  # kill switch on, is_aws_env false
+        ("false", "", False, []),  # kill switch off, is_aws_env false
+    ],
+)
+def test_add_execution_tag_lambda_not_traced(
+    kill_switch_value,
+    is_aws_environment_value,
+    expected_ret_value,
+    expected_tags,
+    capsys,
+    monkeypatch,
+):
+    monkeypatch.setenv("LUMIGO_SWITCH_OFF", kill_switch_value)
+    monkeypatch.setenv("AWS_LAMBDA_FUNCTION_VERSION", is_aws_environment_value)
+
+    assert add_execution_tag("key", "my-value") is expected_ret_value
+    if expected_ret_value is False:
+        assert "Unable to add tag" in capsys.readouterr().out
+    assert SpansContainer.get_span().execution_tags == expected_tags

src/test/unit/wrappers/http/test_sync_http_wrappers.py

@@ -344,7 +344,7 @@ def lambda_test_function(event, context):
     assert span["info"]["httpInfo"]["request"]["method"] == "POST"
 
 
-def test_wrapping_with_tags_for_api_gw_headers(monkeypatch, context, token):
+def test_wrapping_with_tags_for_api_gw_headers(monkeypatch, context, token, lambda_traced):
     monkeypatch.setattr(auto_tag_event, "AUTO_TAG_API_GW_HEADERS", ["Accept"])
 
     @lumigo_tracer.lumigo_tracer()