Skip to content

Commit c8bd364

Browse files
michael-zhao459michael-zhao459
committed
add tests
1 parent 1a94b04 commit c8bd364

File tree

3 files changed

+416
-1
lines changed

3 files changed

+416
-1
lines changed

‎datadog_lambda/tracing.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -268,6 +268,7 @@ def extract_context_from_sqs_or_sns_event_or_context(event, lambda_context):
268268
"Datadog Lambda Python only supports extracting trace"
269269
"context from String or Binary SQS/SNS message attributes"
270270
)
271+
271272
if dd_json_data:
272273
dd_data = json.loads(dd_json_data)
273274

‎tests/test_tracing.py‎

Lines changed: 177 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import pytest
66
import os
77
import unittest
8+
import base64
89

910
from unittest.mock import Mock, patch, call
1011

@@ -41,6 +42,8 @@
4142
service_mapping as global_service_mapping,
4243
propagator,
4344
emit_telemetry_on_exception_outside_of_handler,
45+
extract_context_from_sqs_or_sns_event_or_context,
46+
extract_context_from_kinesis_event,
4447
)
4548

4649
from tests.utils import get_mock_context
@@ -714,7 +717,7 @@ def test_step_function_trace_data_redrive(self):
714717

715718
@with_trace_propagation_style("datadog")
716719
def test_step_function_trace_data_lambda_root(self):
717-
"""Test JSONata style step function trace data extraction where there's an upstream Lambda"""
720+
"""Test JSONata style step function trace data extraction with upstream Lambda."""
718721
sfn_event = {
719722
"_datadog": {
720723
"Execution": {
@@ -2438,3 +2441,176 @@ def test_exception_outside_handler_tracing_disabled(
24382441

24392442
mock_submit_errors_metric.assert_called_once_with(None)
24402443
mock_trace.assert_not_called()
2444+
2445+
2446+
class TestExtractContextHaveAddedReturnValues(unittest.TestCase):
2447+
def test_extract_context_from_sqs_event_with_datadog_context(self):
2448+
"""Test SQS event extraction with datadog context returns expected values."""
2449+
lambda_ctx = get_mock_context()
2450+
expected_trace_data = {
2451+
TraceHeader.TRACE_ID: "456",
2452+
TraceHeader.PARENT_ID: "789",
2453+
TraceHeader.SAMPLING_PRIORITY: "1",
2454+
}
2455+
sqs_event = {
2456+
"Records": [
2457+
{
2458+
"messageId": "test-message-id",
2459+
"eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue",
2460+
"messageAttributes": {
2461+
"_datadog": {
2462+
"stringValue": json.dumps(expected_trace_data),
2463+
"dataType": "String",
2464+
}
2465+
},
2466+
}
2467+
]
2468+
}
2469+
2470+
(context, dd_json_data, arn) = extract_context_from_sqs_or_sns_event_or_context(
2471+
sqs_event, lambda_ctx
2472+
)
2473+
2474+
self.assertEqual(context.trace_id, 456)
2475+
self.assertEqual(context.span_id, 789)
2476+
self.assertEqual(context.sampling_priority, 1)
2477+
self.assertEqual(dd_json_data, expected_trace_data)
2478+
self.assertEqual(arn, "arn:aws:sqs:us-east-1:123456789012:test-queue")
2479+
2480+
def test_extract_context_from_sns_event_with_datadog_context(self):
2481+
"""Test SNS event extraction with datadog context returns expected values."""
2482+
lambda_ctx = get_mock_context()
2483+
expected_trace_data = {
2484+
TraceHeader.TRACE_ID: "111",
2485+
TraceHeader.PARENT_ID: "222",
2486+
TraceHeader.SAMPLING_PRIORITY: "2",
2487+
}
2488+
sns_event = {
2489+
"Records": [
2490+
{
2491+
"Sns": {
2492+
"TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic",
2493+
"MessageAttributes": {
2494+
"_datadog": {
2495+
"Type": "String",
2496+
"Value": json.dumps(expected_trace_data),
2497+
}
2498+
},
2499+
}
2500+
}
2501+
]
2502+
}
2503+
2504+
(context, dd_json_data, arn) = extract_context_from_sqs_or_sns_event_or_context(
2505+
sns_event, lambda_ctx
2506+
)
2507+
2508+
self.assertEqual(context.trace_id, 111)
2509+
self.assertEqual(context.span_id, 222)
2510+
self.assertEqual(context.sampling_priority, 2)
2511+
self.assertEqual(dd_json_data, expected_trace_data)
2512+
self.assertEqual(arn, "arn:aws:sns:us-east-1:123456789012:test-topic")
2513+
2514+
def test_extract_context_from_kinesis_event_with_datadog_context(self):
2515+
"""Test Kinesis event extraction with datadog context returns expected values."""
2516+
lambda_ctx = get_mock_context()
2517+
expected_trace_data = {
2518+
TraceHeader.TRACE_ID: "333",
2519+
TraceHeader.PARENT_ID: "444",
2520+
}
2521+
kinesis_event = {
2522+
"Records": [
2523+
{
2524+
"eventSourceARN": "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream",
2525+
"kinesis": {
2526+
"data": base64.b64encode(
2527+
json.dumps({"_datadog": expected_trace_data}).encode()
2528+
).decode()
2529+
},
2530+
}
2531+
]
2532+
}
2533+
2534+
context, dd_json_data, arn = extract_context_from_kinesis_event(
2535+
kinesis_event, lambda_ctx
2536+
)
2537+
2538+
self.assertEqual(context.trace_id, 333)
2539+
self.assertEqual(context.span_id, 444)
2540+
self.assertEqual(dd_json_data, expected_trace_data)
2541+
self.assertEqual(
2542+
arn, "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream"
2543+
)
2544+
2545+
def test_step_function_context_in_sqs_calls_step_function_extractor(self):
2546+
"""Test that Step Function context in SQS calls the step function extractor."""
2547+
lambda_ctx = get_mock_context()
2548+
sfn_context = {
2549+
"Execution": {"Id": "test-execution-id"},
2550+
"StateMachine": {"Id": "test-state-machine-id"},
2551+
"State": {"Name": "test-state"},
2552+
}
2553+
sfn_sqs_event = {
2554+
"Records": [
2555+
{
2556+
"eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue",
2557+
"messageAttributes": {
2558+
"_datadog": {
2559+
"stringValue": json.dumps(sfn_context),
2560+
"dataType": "String",
2561+
}
2562+
},
2563+
}
2564+
]
2565+
}
2566+
2567+
(context, dd_json_data, arn) = extract_context_from_sqs_or_sns_event_or_context(
2568+
sfn_sqs_event, lambda_ctx
2569+
)
2570+
2571+
self.assertIsNotNone(context)
2572+
self.assertIsNone(dd_json_data)
2573+
self.assertIsNone(arn)
2574+
2575+
def test_extract_context_from_sns_to_sqs_event_with_datadog_context(self):
2576+
"""Test SNS -> SQS event extraction with datadog context returns expected values."""
2577+
lambda_ctx = get_mock_context()
2578+
expected_trace_data = {
2579+
TraceHeader.TRACE_ID: "555",
2580+
TraceHeader.PARENT_ID: "666",
2581+
TraceHeader.SAMPLING_PRIORITY: "1",
2582+
}
2583+
2584+
sns_body = {
2585+
"Type": "Notification",
2586+
"TopicArn": "arn:aws:sns:us-east-1:123456789012:test-topic",
2587+
"MessageAttributes": {
2588+
"_datadog": {
2589+
"Type": "String",
2590+
"Value": json.dumps(expected_trace_data),
2591+
}
2592+
},
2593+
"Message": "Test message",
2594+
"MessageId": "test-message-id",
2595+
"Timestamp": "2023-01-01T00:00:00.000Z",
2596+
}
2597+
2598+
sns_to_sqs_event = {
2599+
"Records": [
2600+
{
2601+
"messageId": "sqs-message-id",
2602+
"eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:test-queue",
2603+
"body": json.dumps(sns_body),
2604+
}
2605+
]
2606+
}
2607+
2608+
(context, dd_json_data, arn) = extract_context_from_sqs_or_sns_event_or_context(
2609+
sns_to_sqs_event, lambda_ctx
2610+
)
2611+
2612+
self.assertEqual(context.trace_id, 555)
2613+
self.assertEqual(context.span_id, 666)
2614+
self.assertEqual(context.sampling_priority, 1)
2615+
self.assertEqual(dd_json_data, expected_trace_data)
2616+
self.assertEqual(arn, "arn:aws:sqs:us-east-1:123456789012:test-queue")

0 commit comments

Comments
 (0)