diff --git a/datadog_lambda/tracing.py b/datadog_lambda/tracing.py index 7879697e..2834f977 100644 --- a/datadog_lambda/tracing.py +++ b/datadog_lambda/tracing.py @@ -685,9 +685,36 @@ def create_inferred_span( return None +def create_service_mapping(val): + new_service_mapping = {} + for entry in val.split(","): + parts = entry.split(":") + if len(parts) == 2: + key = parts[0].strip() + value = parts[1].strip() + if key != value and key and value: + new_service_mapping[key] = value + return new_service_mapping + + +def determine_service_name(service_mapping, specific_key, generic_key, default_value): + service_name = service_mapping.get(specific_key) + if service_name is None: + service_name = service_mapping.get(generic_key, default_value) + return service_name + + +service_mapping = {} +# Initialization code +service_mapping_str = os.getenv("DD_SERVICE_MAPPING", "") +service_mapping = create_service_mapping(service_mapping_str) + + def create_inferred_span_from_lambda_function_url_event(event, context): request_context = event.get("requestContext") + api_id = request_context.get("apiId") domain = request_context.get("domainName") + service_name = determine_service_name(service_mapping, api_id, "lambda_url", domain) method = request_context.get("http", {}).get("method") path = request_context.get("http", {}).get("path") resource = "{0} {1}".format(method, path) @@ -701,7 +728,7 @@ def create_inferred_span_from_lambda_function_url_event(event, context): } request_time_epoch = request_context.get("timeEpoch") args = { - "service": domain, + "service": service_name, "resource": resource, "span_type": "http", } @@ -790,13 +817,18 @@ def create_inferred_span_from_api_gateway_websocket_event( request_context = event.get("requestContext") domain = request_context.get("domainName") endpoint = request_context.get("routeKey") + api_id = request_context.get("apiId") + + service_name = determine_service_name( + service_mapping, api_id, "lambda_api_gateway", domain + ) tags = { "operation_name": "aws.apigateway.websocket", "http.url": domain + endpoint, "endpoint": endpoint, "resource_names": endpoint, - "apiid": request_context.get("apiId"), - "apiname": request_context.get("apiId"), + "apiid": api_id, + "apiname": api_id, "stage": request_context.get("stage"), "request_id": context.aws_request_id, "connection_id": request_context.get("connectionId"), @@ -809,7 +841,7 @@ def create_inferred_span_from_api_gateway_websocket_event( else: InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="sync") args = { - "service": domain, + "service": service_name, "resource": endpoint, "span_type": "web", } @@ -838,6 +870,10 @@ def create_inferred_span_from_api_gateway_event( ): request_context = event.get("requestContext") domain = request_context.get("domainName", "") + api_id = request_context.get("apiId") + service_name = determine_service_name( + service_mapping, api_id, "lambda_api_gateway", domain + ) method = event.get("httpMethod") path = event.get("path") resource = "{0} {1}".format(method, path) @@ -847,8 +883,8 @@ def create_inferred_span_from_api_gateway_event( "endpoint": path, "http.method": method, "resource_names": resource, - "apiid": request_context.get("apiId"), - "apiname": request_context.get("apiId"), + "apiid": api_id, + "apiname": api_id, "stage": request_context.get("stage"), "request_id": context.aws_request_id, } @@ -858,7 +894,7 @@ def create_inferred_span_from_api_gateway_event( else: InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="sync") args = { - "service": domain, + "service": service_name, "resource": resource, "span_type": "http", } @@ -888,6 +924,10 @@ def create_inferred_span_from_http_api_event( ): request_context = event.get("requestContext") domain = request_context.get("domainName") + api_id = request_context.get("apiId") + service_name = determine_service_name( + service_mapping, api_id, "lambda_api_gateway", domain + ) method = request_context.get("http", {}).get("method") path = event.get("rawPath") resource = "{0} {1}".format(method, path) @@ -901,8 +941,8 @@ def create_inferred_span_from_http_api_event( "http.user_agent": request_context.get("http", {}).get("userAgent"), "resource_names": resource, "request_id": context.aws_request_id, - "apiid": request_context.get("apiId"), - "apiname": request_context.get("apiId"), + "apiid": api_id, + "apiname": api_id, "stage": request_context.get("stage"), } request_time_epoch_ms = int(request_context.get("timeEpoch")) @@ -911,7 +951,7 @@ def create_inferred_span_from_http_api_event( else: InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="sync") args = { - "service": domain, + "service": service_name, "resource": resource, "span_type": "http", } @@ -936,6 +976,9 @@ def create_inferred_span_from_sqs_event(event, context): event_record = get_first_record(event) event_source_arn = event_record.get("eventSourceARN") queue_name = event_source_arn.split(":")[-1] + service_name = determine_service_name( + service_mapping, queue_name, "lambda_sqs", "sqs" + ) tags = { "operation_name": "aws.sqs", "resource_names": queue_name, @@ -947,7 +990,7 @@ def create_inferred_span_from_sqs_event(event, context): InferredSpanInfo.set_tags(tags, tag_source="self", synchronicity="async") request_time_epoch = event_record.get("attributes", {}).get("SentTimestamp") args = { - "service": "sqs", + "service": service_name, "resource": queue_name, "span_type": "web", } @@ -990,6 +1033,9 @@ def create_inferred_span_from_sns_event(event, context): sns_message = event_record.get("Sns") topic_arn = event_record.get("Sns", {}).get("TopicArn") topic_name = topic_arn.split(":")[-1] + service_name = determine_service_name( + service_mapping, topic_name, "lambda_sns", "sns" + ) tags = { "operation_name": "aws.sns", "resource_names": topic_name, @@ -1009,7 +1055,7 @@ def create_inferred_span_from_sns_event(event, context): dt = datetime.strptime(timestamp, sns_dt_format) args = { - "service": "sns", + "service": service_name, "resource": topic_name, "span_type": "web", } @@ -1027,6 +1073,9 @@ def create_inferred_span_from_kinesis_event(event, context): event_id = event_record.get("eventID") stream_name = event_source_arn.split(":")[-1] shard_id = event_id.split(":")[0] + service_name = determine_service_name( + service_mapping, stream_name, "lambda_kinesis", "kinesis" + ) tags = { "operation_name": "aws.kinesis", "resource_names": stream_name, @@ -1044,7 +1093,7 @@ def create_inferred_span_from_kinesis_event(event, context): ) args = { - "service": "kinesis", + "service": service_name, "resource": stream_name, "span_type": "web", } @@ -1060,6 +1109,9 @@ def create_inferred_span_from_dynamodb_event(event, context): event_record = get_first_record(event) event_source_arn = event_record.get("eventSourceARN") table_name = event_source_arn.split("/")[1] + service_name = determine_service_name( + service_mapping, table_name, "lambda_dynamodb", "dynamodb" + ) dynamodb_message = event_record.get("dynamodb") tags = { "operation_name": "aws.dynamodb", @@ -1077,7 +1129,7 @@ def create_inferred_span_from_dynamodb_event(event, context): "ApproximateCreationDateTime" ) args = { - "service": "dynamodb", + "service": service_name, "resource": table_name, "span_type": "web", } @@ -1093,6 +1145,9 @@ def create_inferred_span_from_dynamodb_event(event, context): def create_inferred_span_from_s3_event(event, context): event_record = get_first_record(event) bucket_name = event_record.get("s3", {}).get("bucket", {}).get("name") + service_name = determine_service_name( + service_mapping, bucket_name, "lambda_s3", "s3" + ) tags = { "operation_name": "aws.s3", "resource_names": bucket_name, @@ -1109,7 +1164,7 @@ def create_inferred_span_from_s3_event(event, context): dt = datetime.strptime(timestamp, dt_format) args = { - "service": "s3", + "service": service_name, "resource": bucket_name, "span_type": "web", } @@ -1123,6 +1178,9 @@ def create_inferred_span_from_s3_event(event, context): def create_inferred_span_from_eventbridge_event(event, context): source = event.get("source") + service_name = determine_service_name( + service_mapping, source, "lambda_eventbridge", "eventbridge" + ) tags = { "operation_name": "aws.eventbridge", "resource_names": source, @@ -1138,7 +1196,7 @@ def create_inferred_span_from_eventbridge_event(event, context): dt = datetime.strptime(timestamp, dt_format) args = { - "service": "eventbridge", + "service": service_name, "resource": source, "span_type": "web", } diff --git a/tests/event_samples/lambda-url.json b/tests/event_samples/lambda-url.json new file mode 100644 index 00000000..8e78650a --- /dev/null +++ b/tests/event_samples/lambda-url.json @@ -0,0 +1,47 @@ +{ + "version": "2.0", + "routeKey": "$default", + "rawPath": "/", + "rawQueryString": "", + "headers": { + "sec-fetch-mode": "navigate", + "sec-fetch-site": "none", + "accept-language": "en-US,en;q=0.9", + "x-forwarded-proto": "https", + "x-forwarded-port": "443", + "x-forwarded-for": "71.195.30.42", + "sec-fetch-user": "?1", + "pragma": "no-cache", + "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.9", + "sec-ch-ua": "\"Google Chrome\";v=\"95\", \"Chromium\";v=\"95\", \";Not A Brand\";v=\"99\"", + "sec-ch-ua-mobile": "?0", + "x-amzn-trace-id": "Root=1-61953929-1ec00c3011062a48477b169e", + "sec-ch-ua-platform": "\"macOS\"", + "host": "a8hyhsshac.lambda-url.eu-south-1.amazonaws.com", + "upgrade-insecure-requests": "1", + "cache-control": "no-cache", + "accept-encoding": "gzip, deflate, br", + "sec-fetch-dest": "document", + "user-agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/95.0.4638.69 Safari/537.36" + }, + "requestContext": { + "accountId": "601427279990", + "apiId": "a8hyhsshac", + "domainName": "a8hyhsshac.lambda-url.eu-south-1.amazonaws.com", + "domainPrefix": "a8hyhsshac", + "http": { + "method": "GET", + "path": "/", + "protocol": "HTTP/1.1", + "sourceIp": "71.195.30.42", + "userAgent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/95.0.4638.69 Safari/537.36" + }, + "requestId": "ec4d58f8-2b8b-4ceb-a1d5-2be7bff58505", + "routeKey": "$default", + "stage": "$default", + "time": "17/Nov/2021:17:17:29 +0000", + "timeEpoch": 1637169449721 + }, + "isBase64Encoded": false + } + \ No newline at end of file diff --git a/tests/test_tracing.py b/tests/test_tracing.py index e19c66aa..01ef3e20 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -1,6 +1,7 @@ import unittest import json import os +import copy from unittest.mock import MagicMock, Mock, patch, call @@ -29,6 +30,9 @@ _convert_xray_sampling, InferredSpanInfo, extract_context_from_eventbridge_event, + create_service_mapping, + determine_service_name, + service_mapping as global_service_mapping, ) from datadog_lambda.trigger import EventTypes @@ -732,6 +736,509 @@ def _basic_common_checks( self.assertEqual(span.get_tag("request_id"), "abc123") +class TestServiceMapping(unittest.TestCase): + def setUp(self): + self.service_mapping = {} + + def get_service_mapping(self): + return global_service_mapping + + def set_service_mapping(self, new_service_mapping): + global_service_mapping.clear() + global_service_mapping.update(new_service_mapping) + + def test_create_service_mapping_invalid_input(self): + # Test case where the input is a string without a colon to split on + val = "api1" + self.assertEqual(create_service_mapping(val), {}) + + # Test case where the input is an empty string + val = "" + self.assertEqual(create_service_mapping(val), {}) + + # Test case where the key and value are identical + val = "api1:api1" + self.assertEqual(create_service_mapping(val), {}) + + # Test case where the key or value is missing + val = ":api1" + self.assertEqual(create_service_mapping(val), {}) + val = "api1:" + self.assertEqual(create_service_mapping(val), {}) + + def test_create_service_mapping(self): + val = "api1:service1,api2:service2" + expected_output = {"api1": "service1", "api2": "service2"} + self.assertEqual(create_service_mapping(val), expected_output) + + def test_get_service_mapping(self): + os.environ["DD_SERVICE_MAPPING"] = "api1:service1,api2:service2" + expected_output = {"api1": "service1", "api2": "service2"} + self.set_service_mapping( + create_service_mapping(os.environ["DD_SERVICE_MAPPING"]) + ) + self.assertEqual(self.get_service_mapping(), expected_output) + + def test_set_service_mapping(self): + new_service_mapping = {"api3": "service3", "api4": "service4"} + self.set_service_mapping(new_service_mapping) + self.assertEqual(self.get_service_mapping(), new_service_mapping) + + def test_determine_service_name(self): + # Prepare the environment + os.environ["DD_SERVICE_MAPPING"] = "api1:service1,api2:service2" + self.set_service_mapping( + create_service_mapping(os.environ["DD_SERVICE_MAPPING"]) + ) + + # Case where specific key is in the service mapping + specific_key = "api1" + self.assertEqual( + determine_service_name( + self.get_service_mapping(), specific_key, "lambda_url", "default" + ), + "service1", + ) + + # Case where specific key is not in the service mapping, but generic key is + specific_key = "api3" + self.assertEqual( + determine_service_name( + self.get_service_mapping(), specific_key, "api2", "default" + ), + "service2", + ) + + # Case where neither specific nor generic keys are in the service mapping + specific_key = "api3" + self.assertEqual( + determine_service_name( + self.get_service_mapping(), specific_key, "api3", "default" + ), + "default", + ) + + def test_remaps_all_inferred_span_service_names_from_api_gateway_event(self): + new_service_mapping = {"lambda_api_gateway": "new-name"} + self.set_service_mapping(new_service_mapping) + event_sample_source = "api-gateway" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + original_event = json.load(event) + + ctx = get_mock_context() + ctx.aws_request_id = "123" + + span1 = create_inferred_span(original_event, ctx) + self.assertEqual(span1.get_tag("operation_name"), "aws.apigateway.rest") + self.assertEqual(span1.service, "new-name") + + # Testing the second event + event2 = copy.deepcopy(original_event) + event2["requestContext"][ + "domainName" + ] = "different.execute-api.us-east-2.amazonaws.com" + span2 = create_inferred_span(event2, ctx) + self.assertEqual(span2.get_tag("operation_name"), "aws.apigateway.rest") + self.assertEqual(span2.service, "new-name") + + def test_remaps_specific_inferred_span_service_names_from_api_gateway_event( + self, + ): + new_service_mapping = {"1234567890": "new-name"} + self.set_service_mapping(new_service_mapping) + event_sample_source = "api-gateway" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + original_event = json.load(event) + + ctx = get_mock_context() + ctx.aws_request_id = "123" + + span1 = create_inferred_span(original_event, ctx) + self.assertEqual(span1.get_tag("operation_name"), "aws.apigateway.rest") + self.assertEqual(span1.service, "new-name") + + # Testing the second event + event2 = copy.deepcopy(original_event) + event2["requestContext"]["apiId"] = "different" + span2 = create_inferred_span(event2, ctx) + self.assertEqual(span2.get_tag("operation_name"), "aws.apigateway.rest") + self.assertEqual( + span2.service, "70ixmpl4fl.execute-api.us-east-2.amazonaws.com" + ) + + def test_remaps_specific_inferred_span_service_names_from_api_gateway_websocket_event( + self, + ): + self.set_service_mapping({"p62c47itsb": "new-name"}) + event_sample_source = "api-gateway-websocket-default" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + original_event = json.load(event) + + ctx = get_mock_context() + ctx.aws_request_id = "123" + + span1 = create_inferred_span(original_event, ctx) + self.assertEqual(span1.get_tag("operation_name"), "aws.apigateway.websocket") + self.assertEqual(span1.service, "new-name") + + # Testing the second event + event2 = copy.deepcopy(original_event) + event2["requestContext"]["apiId"] = "different" + span2 = create_inferred_span(event2, ctx) + self.assertEqual(span2.get_tag("operation_name"), "aws.apigateway.websocket") + self.assertEqual( + span2.service, "p62c47itsb.execute-api.eu-west-1.amazonaws.com" + ) + + def test_remaps_specific_inferred_span_service_names_from_api_gateway_http_event( + self, + ): + self.set_service_mapping({"x02yirxc7a": "new-name"}) + event_sample_source = "http-api" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + original_event = json.load(event) + + ctx = get_mock_context() + ctx.aws_request_id = "123" + + span1 = create_inferred_span(original_event, ctx) + self.assertEqual(span1.get_tag("operation_name"), "aws.httpapi") + self.assertEqual(span1.service, "new-name") + + # Testing the second event + event2 = copy.deepcopy(original_event) + event2["requestContext"]["apiId"] = "different" + span2 = create_inferred_span(event2, ctx) + self.assertEqual(span2.get_tag("operation_name"), "aws.httpapi") + self.assertEqual( + span2.service, "x02yirxc7a.execute-api.eu-west-1.amazonaws.com" + ) + + def test_remaps_all_inferred_span_service_names_from_lambda_url_event(self): + self.set_service_mapping({"lambda_url": "new-name"}) + event_sample_source = "lambda-url" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + original_event = json.load(event) + + ctx = get_mock_context() + ctx.aws_request_id = "123" + + span1 = create_inferred_span(original_event, ctx) + self.assertEqual(span1.get_tag("operation_name"), "aws.lambda.url") + self.assertEqual(span1.service, "new-name") + + # Testing the second event + event2 = copy.deepcopy(original_event) + event2["requestContext"][ + "domainName" + ] = "different.lambda-url.eu-south-1.amazonaws.com" + span2 = create_inferred_span(event2, ctx) + self.assertEqual(span2.get_tag("operation_name"), "aws.lambda.url") + self.assertEqual(span2.service, "new-name") + + def test_remaps_specific_inferred_span_service_names_from_lambda_url_event( + self, + ): + self.set_service_mapping({"a8hyhsshac": "new-name"}) + event_sample_source = "lambda-url" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + original_event = json.load(event) + + ctx = get_mock_context() + ctx.aws_request_id = "123" + + span1 = create_inferred_span(original_event, ctx) + self.assertEqual(span1.get_tag("operation_name"), "aws.lambda.url") + self.assertEqual(span1.service, "new-name") + + # Testing the second event + event2 = copy.deepcopy(original_event) + event2["requestContext"]["apiId"] = "different" + span2 = create_inferred_span(event2, ctx) + self.assertEqual(span2.get_tag("operation_name"), "aws.lambda.url") + self.assertEqual( + span2.service, "a8hyhsshac.lambda-url.eu-south-1.amazonaws.com" + ) + + def test_remaps_all_inferred_span_service_names_from_sqs_event(self): + self.set_service_mapping({"lambda_sqs": "new-name"}) + event_sample_source = "sqs-string-msg-attribute" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + original_event = json.load(event) + + ctx = get_mock_context() + ctx.aws_request_id = "123" + + span1 = create_inferred_span(original_event, ctx) + self.assertEqual(span1.get_tag("operation_name"), "aws.sqs") + self.assertEqual(span1.service, "new-name") + + # Testing the second event + event2 = copy.deepcopy(original_event) + event2["Records"][0][ + "eventSourceARN" + ] = "arn:aws:sqs:eu-west-1:123456789012:different-sqs-url" + span2 = create_inferred_span(event2, ctx) + self.assertEqual(span2.get_tag("operation_name"), "aws.sqs") + self.assertEqual(span2.service, "new-name") + + def test_remaps_specific_inferred_span_service_names_from_sqs_event(self): + self.set_service_mapping({"InferredSpansQueueNode": "new-name"}) + event_sample_source = "sqs-string-msg-attribute" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + original_event = json.load(event) + + ctx = get_mock_context() + ctx.aws_request_id = "123" + + span1 = create_inferred_span(original_event, ctx) + self.assertEqual(span1.get_tag("operation_name"), "aws.sqs") + self.assertEqual(span1.service, "new-name") + + # Testing the second event + event2 = copy.deepcopy(original_event) + event2["Records"][0][ + "eventSourceARN" + ] = "arn:aws:sqs:eu-west-1:123456789012:different-sqs-url" + span2 = create_inferred_span(event2, ctx) + self.assertEqual(span2.get_tag("operation_name"), "aws.sqs") + self.assertEqual(span2.service, "sqs") + + def test_remaps_all_inferred_span_service_names_from_sns_event(self): + self.set_service_mapping({"lambda_sns": "new-name"}) + event_sample_source = "sns-string-msg-attribute" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + original_event = json.load(event) + + ctx = get_mock_context() + ctx.aws_request_id = "123" + + span1 = create_inferred_span(original_event, ctx) + self.assertEqual(span1.get_tag("operation_name"), "aws.sns") + self.assertEqual(span1.service, "new-name") + + # Testing the second event + event2 = copy.deepcopy(original_event) + event2["Records"][0]["Sns"][ + "TopicArn" + ] = "arn:aws:sns:us-west-2:123456789012:different-sns-topic" + span2 = create_inferred_span(event2, ctx) + self.assertEqual(span2.get_tag("operation_name"), "aws.sns") + self.assertEqual(span2.service, "new-name") + + def test_remaps_specific_inferred_span_service_names_from_sns_event(self): + self.set_service_mapping({"serverlessTracingTopicPy": "new-name"}) + event_sample_source = "sns-string-msg-attribute" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + original_event = json.load(event) + + ctx = get_mock_context() + ctx.aws_request_id = "123" + + span1 = create_inferred_span(original_event, ctx) + self.assertEqual(span1.get_tag("operation_name"), "aws.sns") + self.assertEqual(span1.service, "new-name") + + # Testing the second event + event2 = copy.deepcopy(original_event) + event2["Records"][0]["Sns"][ + "TopicArn" + ] = "arn:aws:sns:us-west-2:123456789012:different-sns-topic" + span2 = create_inferred_span(event2, ctx) + self.assertEqual(span2.get_tag("operation_name"), "aws.sns") + self.assertEqual(span2.service, "sns") + + def test_remaps_all_inferred_span_service_names_from_kinesis_event(self): + self.set_service_mapping({"lambda_kinesis": "new-name"}) + event_sample_source = "kinesis" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + original_event = json.load(event) + + ctx = get_mock_context() + ctx.aws_request_id = "123" + + span1 = create_inferred_span(original_event, ctx) + self.assertEqual(span1.get_tag("operation_name"), "aws.kinesis") + self.assertEqual(span1.service, "new-name") + + # Testing the second event + event2 = copy.deepcopy(original_event) + event2["Records"][0][ + "eventSourceARN" + ] = "arn:aws:kinesis:eu-west-1:601427279990:stream/differentKinesisStream" + span2 = create_inferred_span(event2, ctx) + self.assertEqual(span2.get_tag("operation_name"), "aws.kinesis") + self.assertEqual(span2.service, "new-name") + + def test_remaps_specific_inferred_span_service_names_from_kinesis_event(self): + self.set_service_mapping({"Different_EXAMPLE": "new-name"}) + event_sample_source = "kinesis" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + original_event = json.load(event) + + ctx = get_mock_context() + ctx.aws_request_id = "123" + + span1 = create_inferred_span(original_event, ctx) + self.assertEqual(span1.get_tag("operation_name"), "aws.kinesis") + self.assertEqual(span1.service, "kinesis") + + # Testing the second event + event2 = copy.deepcopy(original_event) + event2["Records"][0][ + "eventSourceARN" + ] = "arn:aws:kinesis:eu-west-1:601427279990:stream/DifferentKinesisStream" + span2 = create_inferred_span(event2, ctx) + self.assertEqual(span2.get_tag("operation_name"), "aws.kinesis") + self.assertEqual(span2.service, "kinesis") + + def test_remaps_all_inferred_span_service_names_from_s3_event(self): + self.set_service_mapping({"lambda_s3": "new-name"}) + event_sample_source = "s3" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + original_event = json.load(event) + + ctx = get_mock_context() + ctx.aws_request_id = "123" + + span1 = create_inferred_span(original_event, ctx) + self.assertEqual(span1.get_tag("operation_name"), "aws.s3") + self.assertEqual(span1.service, "new-name") + + # Testing the second event + event2 = copy.deepcopy(original_event) + event2["Records"][0]["s3"]["bucket"][ + "arn" + ] = "arn:aws:s3:::different-example-bucket" + span2 = create_inferred_span(event2, ctx) + self.assertEqual(span2.get_tag("operation_name"), "aws.s3") + self.assertEqual(span2.service, "new-name") + + def test_remaps_specific_inferred_span_service_names_from_s3_event(self): + self.set_service_mapping({"example-bucket": "new-name"}) + event_sample_source = "s3" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + original_event = json.load(event) + + ctx = get_mock_context() + ctx.aws_request_id = "123" + + span1 = create_inferred_span(original_event, ctx) + self.assertEqual(span1.get_tag("operation_name"), "aws.s3") + self.assertEqual(span1.service, "new-name") + + # Testing the second event + event2 = copy.deepcopy(original_event) + event2["Records"][0]["s3"]["bucket"]["name"] = "different-example-bucket" + span2 = create_inferred_span(event2, ctx) + self.assertEqual(span2.get_tag("operation_name"), "aws.s3") + self.assertEqual(span2.service, "s3") + + def test_remaps_all_inferred_span_service_names_from_dynamodb_event(self): + self.set_service_mapping({"lambda_dynamodb": "new-name"}) + event_sample_source = "dynamodb" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + original_event = json.load(event) + + ctx = get_mock_context() + ctx.aws_request_id = "123" + + span1 = create_inferred_span(original_event, ctx) + self.assertEqual(span1.get_tag("operation_name"), "aws.dynamodb") + self.assertEqual(span1.service, "new-name") + + # Testing the second event + event2 = copy.deepcopy(original_event) + event2["Records"][0][ + "eventSourceARN" + ] = "arn:aws:dynamodb:us-east-1:123456789012:table/DifferentExampleTableWithStream/stream/2015-06-27T00:48:05.899" + span2 = create_inferred_span(event2, ctx) + self.assertEqual(span2.get_tag("operation_name"), "aws.dynamodb") + self.assertEqual(span2.service, "new-name") + + def test_remaps_specific_inferred_span_service_names_from_dynamodb_event(self): + self.set_service_mapping({"ExampleTableWithStream": "new-name"}) + event_sample_source = "dynamodb" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + original_event = json.load(event) + + ctx = get_mock_context() + ctx.aws_request_id = "123" + + span1 = create_inferred_span(original_event, ctx) + self.assertEqual(span1.get_tag("operation_name"), "aws.dynamodb") + self.assertEqual(span1.service, "new-name") + + # Testing the second event + event2 = copy.deepcopy(original_event) + event2["Records"][0][ + "eventSourceARN" + ] = "arn:aws:dynamodb:us-east-1:123456789012:table/DifferentExampleTableWithStream/stream/2015-06-27T00:48:05.899" + span2 = create_inferred_span(event2, ctx) + self.assertEqual(span2.get_tag("operation_name"), "aws.dynamodb") + self.assertEqual(span2.service, "dynamodb") + + def test_remaps_all_inferred_span_service_names_from_eventbridge_event(self): + self.set_service_mapping({"lambda_eventbridge": "new-name"}) + event_sample_source = "eventbridge-custom" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + original_event = json.load(event) + + ctx = get_mock_context() + ctx.aws_request_id = "123" + + span1 = create_inferred_span(original_event, ctx) + self.assertEqual(span1.get_tag("operation_name"), "aws.eventbridge") + self.assertEqual(span1.service, "new-name") + + # Testing the second event + event2 = copy.deepcopy(original_event) + event2["source"] = "different.eventbridge.custom.event.sender" + span2 = create_inferred_span(event2, ctx) + self.assertEqual(span2.get_tag("operation_name"), "aws.eventbridge") + self.assertEqual(span2.service, "new-name") + + def test_remaps_specific_inferred_span_service_names_from_eventbridge_event( + self, + ): + self.set_service_mapping({"eventbridge.custom.event.sender": "new-name"}) + event_sample_source = "eventbridge-custom" + test_file = event_samples + event_sample_source + ".json" + with open(test_file, "r") as event: + original_event = json.load(event) + + ctx = get_mock_context() + ctx.aws_request_id = "123" + + span1 = create_inferred_span(original_event, ctx) + self.assertEqual(span1.get_tag("operation_name"), "aws.eventbridge") + self.assertEqual(span1.service, "new-name") + + # Testing the second event + event2 = copy.deepcopy(original_event) + event2["source"] = "different.eventbridge.custom.event.sender" + span2 = create_inferred_span(event2, ctx) + self.assertEqual(span2.get_tag("operation_name"), "aws.eventbridge") + self.assertEqual(span2.service, "eventbridge") + + class TestInferredSpans(unittest.TestCase): def test_create_inferred_span_from_api_gateway_event(self): event_sample_source = "api-gateway"