diff --git a/datadog_lambda/xray.py b/datadog_lambda/xray.py index d4836fae..2e3a2fe2 100644 --- a/datadog_lambda/xray.py +++ b/datadog_lambda/xray.py @@ -10,11 +10,11 @@ logger = logging.getLogger(__name__) -def get_xray_host_port(adress): - if adress == "": +def get_xray_host_port(address): + if address == "": logger.debug("X-Ray daemon env var not set, not sending sub-segment") return None - parts = adress.split(":") + parts = address.split(":") if len(parts) <= 1: logger.debug("X-Ray daemon env var not set, not sending sub-segment") return None @@ -108,6 +108,11 @@ def send_segment(key, metadata): "Failed to create segment since it was not possible to get trace context from header" ) return None + + # Skip adding segment, if the xray trace is going to be sampled away. + if context["sampled"] == "0": + logger.debug("Skipping sending metadata, x-ray trace was sampled out") + return None segment = build_segment(context, key, metadata) segment_payload = build_segment_payload(segment) send(host_port_tuple, segment_payload) diff --git a/tests/Dockerfile b/tests/Dockerfile index 86bafa0b..948de1bc 100644 --- a/tests/Dockerfile +++ b/tests/Dockerfile @@ -7,7 +7,8 @@ RUN mkdir -p /test/datadog_lambda WORKDIR /test # Copy minimal subset of files to make pip install succeed and be cached (next docker builds will be way faster) -COPY setup.py . +COPY pyproject.toml . +COPY poetry.lock . COPY README.md . COPY datadog_lambda/__init__.py datadog_lambda/__init__.py diff --git a/tests/test_xray.py b/tests/test_xray.py index 5c3cd404..d996f8e8 100644 --- a/tests/test_xray.py +++ b/tests/test_xray.py @@ -1,12 +1,25 @@ import unittest import json +import os from unittest.mock import MagicMock, patch -from datadog_lambda.xray import get_xray_host_port, build_segment_payload, build_segment +from datadog_lambda.xray import ( + get_xray_host_port, + build_segment_payload, + build_segment, + send_segment, +) class TestXRay(unittest.TestCase): + def tearDown(self): + if os.environ.get("_X_AMZN_TRACE_ID"): + os.environ.pop("_X_AMZN_TRACE_ID") + if os.environ.get("AWS_XRAY_DAEMON_ADDRESS"): + os.environ.pop("AWS_XRAY_DAEMON_ADDRESS") + return super().tearDown() + def test_get_xray_host_port_empty_(self): result = get_xray_host_port("") self.assertIsNone(result) @@ -20,6 +33,31 @@ def test_get_xray_host_port_success(self): self.assertEqual("mySuperHost", result[0]) self.assertEqual(1000, result[1]) + def test_send_segment_sampled_out(self): + os.environ["AWS_XRAY_DAEMON_ADDRESS"] = "fake-agent.com:8080" + os.environ[ + "_X_AMZN_TRACE_ID" + ] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=0" + + with patch( + "datadog_lambda.xray.send", MagicMock(return_value=None) + ) as mock_send: + # XRay trace won't be sampled according to the trace header. + send_segment("my_key", {"data": "value"}) + self.assertFalse(mock_send.called) + + def test_send_segment_sampled(self): + os.environ["AWS_XRAY_DAEMON_ADDRESS"] = "fake-agent.com:8080" + os.environ[ + "_X_AMZN_TRACE_ID" + ] = "Root=1-5e272390-8c398be037738dc042009320;Parent=94ae789b969f1cc5;Sampled=1" + with patch( + "datadog_lambda.xray.send", MagicMock(return_value=None) + ) as mock_send: + # X-Ray trace will be sampled according to the trace header. + send_segment("my_key", {"data": "value"}) + self.assertTrue(mock_send.called) + def test_build_segment_payload_ok(self): exected_text = '{"format": "json", "version": 1}\nmyPayload' self.assertEqual(exected_text, build_segment_payload("myPayload"))