Improve testing for ddtrace imports. (#662)

Author: purple4reina

.github/workflows/build.yml

@@ -5,6 +5,8 @@ on:
     branches:
       - "main"
   pull_request:
+  schedule:
+    - cron: '0 0,12 * * *' # Runs every day at midnight and noon utc
 
 jobs:
   lint:

datadog_lambda/wrapper.py

@@ -46,6 +46,10 @@
     extract_http_status_code_tag,
 )
 
+# ddtrace imports are also tested in
+# dd-trace-py/tests/internal/test_serverless.py please update those tests when
+# making changes to any ddtrace import.
+
 if config.appsec_enabled:
     from datadog_lambda.asm import (
         asm_set_context,

scripts/check_layer_size.sh

@@ -8,7 +8,7 @@
 # Compares layer size to threshold, and fails if below that threshold
 
 set -e
-MAX_LAYER_COMPRESSED_SIZE_KB=$(expr 17 \* 1024 / 2)  # 8704 KB
+MAX_LAYER_COMPRESSED_SIZE_KB=$(expr 9 \* 1024)      # 9216 KB
 MAX_LAYER_UNCOMPRESSED_SIZE_KB=$(expr 25 \* 1024)   # 25600 KB
 
 

tests/test_wrapper.py

@@ -789,3 +789,113 @@ def lambda_handler(event, context):
         self.mock_asm_start_response.assert_not_called()
 
         assert lambda_handler.span.get_tag("http.status_code") == "200"
+
+
+@patch("datadog_lambda.config.Config.exception_replay_enabled", True)
+def test_exception_replay_enabled(monkeypatch):
+    importlib.reload(wrapper)
+
+    original_SpanExceptionHandler_enable = wrapper.SpanExceptionHandler.enable
+    SpanExceptionHandler_enable_calls = []
+
+    def SpanExceptionHandler_enable(*args, **kwargs):
+        SpanExceptionHandler_enable_calls.append((args, kwargs))
+        return original_SpanExceptionHandler_enable(*args, **kwargs)
+
+    original_SignalUploader_periodic = wrapper.SignalUploader.periodic
+    SignalUploader_periodic_calls = []
+
+    def SignalUploader_periodic(*args, **kwargs):
+        SignalUploader_periodic_calls.append((args, kwargs))
+        return original_SignalUploader_periodic(*args, **kwargs)
+
+    monkeypatch.setattr(
+        "datadog_lambda.wrapper.SpanExceptionHandler.enable",
+        SpanExceptionHandler_enable,
+    )
+    monkeypatch.setattr(
+        "datadog_lambda.wrapper.SignalUploader.periodic", SignalUploader_periodic
+    )
+
+    expected_response = {
+        "statusCode": 200,
+        "body": "This should be returned",
+    }
+
+    @wrapper.datadog_lambda_wrapper
+    def lambda_handler(event, context):
+        return expected_response
+
+    response = lambda_handler({}, get_mock_context())
+
+    assert response == expected_response
+    assert len(SpanExceptionHandler_enable_calls) == 1
+    assert len(SignalUploader_periodic_calls) == 1
+
+
+@patch("datadog_lambda.config.Config.profiling_enabled", True)
+def test_profiling_enabled(monkeypatch):
+    importlib.reload(wrapper)
+
+    original_Profiler_start = wrapper.profiler.Profiler.start
+    Profiler_start_calls = []
+
+    def Profiler_start(*args, **kwargs):
+        Profiler_start_calls.append((args, kwargs))
+        return original_Profiler_start(*args, **kwargs)
+
+    monkeypatch.setattr("datadog_lambda.wrapper.is_new_sandbox", lambda: True)
+    monkeypatch.setattr(
+        "datadog_lambda.wrapper.profiler.Profiler.start", Profiler_start
+    )
+
+    expected_response = {
+        "statusCode": 200,
+        "body": "This should be returned",
+    }
+
+    @wrapper.datadog_lambda_wrapper
+    def lambda_handler(event, context):
+        return expected_response
+
+    response = lambda_handler({}, get_mock_context())
+
+    assert response == expected_response
+    assert len(Profiler_start_calls) == 1
+
+
+@patch("datadog_lambda.config.Config.llmobs_enabled", True)
+def test_llmobs_enabled(monkeypatch):
+    importlib.reload(wrapper)
+
+    original_LLMObs_enable = wrapper.LLMObs.enable
+    LLMObs_enable_calls = []
+
+    def LLMObs_enable(*args, **kwargs):
+        LLMObs_enable_calls.append((args, kwargs))
+        return original_LLMObs_enable(*args, **kwargs)
+
+    original_LLMObs_flush = wrapper.LLMObs.flush
+    LLMObs_flush_calls = []
+
+    def LLMObs_flush(*args, **kwargs):
+        LLMObs_flush_calls.append((args, kwargs))
+        return original_LLMObs_flush(*args, **kwargs)
+
+    monkeypatch.setattr("datadog_lambda.wrapper.LLMObs.enable", LLMObs_enable)
+    monkeypatch.setattr("datadog_lambda.wrapper.LLMObs.flush", LLMObs_flush)
+
+    expected_response = {
+        "statusCode": 200,
+        "body": "This should be returned",
+    }
+
+    @wrapper.datadog_lambda_wrapper
+    def lambda_handler(event, context):
+        return expected_response
+
+    response = lambda_handler({}, get_mock_context())
+
+    assert response == expected_response
+    assert len(LLMObs_enable_calls) == 1
+    assert len(LLMObs_flush_calls) == 1