From 12d892f5133c2e9ca9dd30c3c73f034a5b9252b4 Mon Sep 17 00:00:00 2001
From: Rey Abolofia <rey.abolofia@datadoghq.com>
Date: Fri, 5 Apr 2024 12:57:32 -0700
Subject: [PATCH 1/2] Benchmarks for operations that use json dump/load.

---
 .gitignore               |  2 +
 pyproject.toml           |  5 +++
 tests/test_benchmarks.py | 80 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 87 insertions(+)
 create mode 100644 tests/test_benchmarks.py

diff --git a/.gitignore b/.gitignore
index 3aeeb1fb..ff0272ba 100644
--- a/.gitignore
+++ b/.gitignore
@@ -41,3 +41,5 @@ pip-log.txt
 
 node_modules
 yarn-lock.json
+
+.benchmarks
diff --git a/pyproject.toml b/pyproject.toml
index 3536347f..7d77a952 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -37,6 +37,7 @@ boto3 = { version = "^1.28.0", optional = true }
 typing_extensions = {version = "^4.0", python = "<3.8"}
 requests = { version ="^2.22.0", optional = true }
 pytest = { version= "^8.0.0", optional = true }
+pytest-benchmark = { version = "^4.0", optional = true }
 flake8 = { version = "^5.0.4", optional = true }
 
 
@@ -45,9 +46,13 @@ dev = [
     "boto3",
     "flake8",
     "pytest",
+    "pytest-benchmark",
     "requests",
 ]
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]
 build-backend = "poetry.core.masonry.api"
+
+[tool.pytest.ini_options]
+addopts = "--benchmark-disable --benchmark-autosave"
diff --git a/tests/test_benchmarks.py b/tests/test_benchmarks.py
new file mode 100644
index 00000000..fe495b91
--- /dev/null
+++ b/tests/test_benchmarks.py
@@ -0,0 +1,80 @@
+import builtins
+import json
+import os
+import pytest
+
+import ddtrace
+
+from datadog_lambda import metric
+from datadog_lambda import tag_object
+from datadog_lambda import tracing
+from datadog_lambda import trigger
+from datadog_lambda import xray
+
+from tests.utils import get_mock_context
+
+
+event_samples_dir = "tests/event_samples"
+event_samples = [f[:-5] for f in os.listdir(event_samples_dir) if f.endswith(".json")]
+
+
+def test_metric_write_metric_point_to_stdout(benchmark, monkeypatch):
+    monkeypatch.setattr(builtins, "print", lambda *a, **k: None)
+    benchmark(
+        metric.write_metric_point_to_stdout,
+        "metric_name",
+        1,
+        tags=[
+            "tag1:value1",
+            "tag2:value2",
+            "tag3:value3",
+        ],
+    )
+
+
+@pytest.mark.parametrize("event", event_samples)
+def test_tag_object_tag_object(event, benchmark):
+    with open(f"{event_samples_dir}/{event}.json") as f:
+        event = json.load(f)
+    span = ddtrace.tracer.start_span("test")
+    benchmark(tag_object.tag_object, span, "function.request", event)
+
+
+@pytest.mark.parametrize("event", event_samples)
+def test_tracing_create_inferred_span(event, benchmark):
+    with open(f"{event_samples_dir}/{event}.json") as f:
+        event = json.load(f)
+    context = get_mock_context()
+    benchmark(tracing.create_inferred_span, event, context)
+
+
+@pytest.mark.parametrize("event", event_samples)
+def test_tracing_extract_dd_trace_context(event, benchmark):
+    with open(f"{event_samples_dir}/{event}.json") as f:
+        event = json.load(f)
+    context = get_mock_context()
+    benchmark(tracing.extract_dd_trace_context, event, context)
+
+
+@pytest.mark.parametrize("event", event_samples)
+def test_trigger_parse_event_source(event, benchmark):
+    with open(f"{event_samples_dir}/{event}.json") as f:
+        event = json.load(f)
+    benchmark(trigger.parse_event_source, event)
+
+
+@pytest.mark.parametrize("event", event_samples)
+def test_trigger_extract_trigger_tags(event, benchmark):
+    with open(f"{event_samples_dir}/{event}.json") as f:
+        event = json.load(f)
+    context = get_mock_context()
+    benchmark(trigger.extract_trigger_tags, event, context)
+
+
+def test_xray_send_segment(benchmark):
+    key = {
+        "trace-id": "12345678901234567890123456789012",
+        "parent-id": "1234567890123456",
+        "sampling-priority": "1",
+    }
+    benchmark(xray.send_segment, key, "trace")

From fce40ec0fbdaaaff98dbdc39c0cd79c43650cbe4 Mon Sep 17 00:00:00 2001
From: Rey Abolofia <rey.abolofia@datadoghq.com>
Date: Fri, 5 Apr 2024 13:23:29 -0700
Subject: [PATCH 2/2] Set required envvars for sending xray segments.

---
 tests/test_benchmarks.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/tests/test_benchmarks.py b/tests/test_benchmarks.py
index fe495b91..d430c255 100644
--- a/tests/test_benchmarks.py
+++ b/tests/test_benchmarks.py
@@ -11,6 +11,8 @@
 from datadog_lambda import trigger
 from datadog_lambda import xray
 
+from datadog_lambda.constants import XrayDaemon, XraySubsegment
+
 from tests.utils import get_mock_context
 
 
@@ -71,10 +73,15 @@ def test_trigger_extract_trigger_tags(event, benchmark):
     benchmark(trigger.extract_trigger_tags, event, context)
 
 
-def test_xray_send_segment(benchmark):
+def test_xray_send_segment(benchmark, monkeypatch):
+    monkeypatch.setenv(XrayDaemon.XRAY_DAEMON_ADDRESS, "localhost:9000")
+    monkeypatch.setenv(
+        XrayDaemon.XRAY_TRACE_ID_HEADER_NAME,
+        "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1;Lineage=c6c5b1b9:0",
+    )
     key = {
         "trace-id": "12345678901234567890123456789012",
         "parent-id": "1234567890123456",
         "sampling-priority": "1",
     }
-    benchmark(xray.send_segment, key, "trace")
+    benchmark(xray.send_segment, XraySubsegment.TRACE_KEY, key)