kinesis fix, tests fix

Author: michael-zhao459

datadog_lambda/tracing.py

@@ -426,9 +426,7 @@ def extract_context_from_kinesis_event(event, lambda_context):
     except Exception as e:
         logger.debug("The trace extractor returned with error %s", e)
     # Still want to set a DSM checkpoint even if DSM context not propagated
-    # Do not want to set checkpoint with "" arn
-    if source_arn:
-        _dsm_set_checkpoint(None, "kinesis", source_arn)
+    _dsm_set_checkpoint(None, "kinesis", source_arn)
     return extract_context_from_lambda_context(lambda_context)
 
 

tests/test_tracing.py

@@ -2857,26 +2857,33 @@ def test_sns_to_sqs_event_with_malformed_datadog_message_attributes(
 
     @patch("datadog_lambda.tracing.extract_context_from_lambda_context")
     @patch("datadog_lambda.tracing._dsm_set_checkpoint")
+    @patch("ddtrace.data_streams.set_consume_checkpoint")
     def test_sqs_sns_event_with_exception_accessing_first_record(
-        self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context
+        self,
+        mock_set_consume_checkpoint,
+        mock_dsm_set_checkpoint,
+        mock_extract_from_lambda_context,
     ):
-        event = {"Records": None}
+        event = {"Records": None, "eventSource": "aws:sqs"}
 
         mock_context = Context(trace_id=123, span_id=456)
         mock_extract_from_lambda_context.return_value = mock_context
 
         result = extract_context_from_sqs_or_sns_event_or_context(
             event, self.lambda_context, parse_event_source(event)
         )
-
-        mock_dsm_set_checkpoint.assert_not_called()
+        mock_set_consume_checkpoint.assert_not_called()
         mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context)
         self.assertEqual(result, mock_context)
 
     @patch("datadog_lambda.tracing.extract_context_from_lambda_context")
     @patch("datadog_lambda.tracing._dsm_set_checkpoint")
+    @patch("ddtrace.data_streams.set_consume_checkpoint")
     def test_sqs_event_with_empty_arn(
-        self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context
+        self,
+        mock_set_consume_checkpoint,
+        mock_dsm_set_checkpoint,
+        mock_extract_from_lambda_context,
     ):
         """Test SQS event with empty eventSourceARN"""
         event = {
@@ -2896,14 +2903,19 @@ def test_sqs_event_with_empty_arn(
             event, self.lambda_context, parse_event_source(event)
         )
 
-        mock_dsm_set_checkpoint.assert_not_called()
+        mock_dsm_set_checkpoint.assert_called_once_with(None, "sqs", "")
+        mock_set_consume_checkpoint.assert_not_called()
         mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context)
         self.assertEqual(result, mock_context)
 
     @patch("datadog_lambda.tracing.extract_context_from_lambda_context")
     @patch("datadog_lambda.tracing._dsm_set_checkpoint")
+    @patch("ddtrace.data_streams.set_consume_checkpoint")
     def test_sns_event_with_empty_arn(
-        self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context
+        self,
+        mock_set_consume_checkpoint,
+        mock_dsm_set_checkpoint,
+        mock_extract_from_lambda_context,
     ):
         """Test SNS event with empty TopicArn"""
         event = {
@@ -2925,14 +2937,19 @@ def test_sns_event_with_empty_arn(
             event, self.lambda_context, parse_event_source(event)
         )
 
-        mock_dsm_set_checkpoint.assert_not_called()
+        mock_dsm_set_checkpoint.assert_called_once_with(None, "sns", "")
+        mock_set_consume_checkpoint.assert_not_called()
         mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context)
         self.assertEqual(result, mock_context)
 
     @patch("datadog_lambda.tracing.extract_context_from_lambda_context")
     @patch("datadog_lambda.tracing._dsm_set_checkpoint")
+    @patch("ddtrace.data_streams.set_consume_checkpoint")
     def test_sns_to_sqs_event_with_empty_arn(
-        self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context
+        self,
+        mock_set_consume_checkpoint,
+        mock_dsm_set_checkpoint,
+        mock_extract_from_lambda_context,
     ):
         """Test SNS->SQS event with empty eventSourceARN"""
         sns_notification = {
@@ -2960,7 +2977,8 @@ def test_sns_to_sqs_event_with_empty_arn(
             event, self.lambda_context, parse_event_source(event)
         )
 
-        mock_dsm_set_checkpoint.assert_not_called()
+        mock_dsm_set_checkpoint.assert_called_once_with(None, "sqs", "")
+        mock_set_consume_checkpoint.assert_not_called()
         mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context)
         self.assertEqual(result, mock_context)
 
@@ -3063,8 +3081,12 @@ def test_kinesis_event_with_malformed_data(
 
     @patch("datadog_lambda.tracing.extract_context_from_lambda_context")
     @patch("datadog_lambda.tracing._dsm_set_checkpoint")
+    @patch("ddtrace.data_streams.set_consume_checkpoint")
     def test_kinesis_event_with_empty_arn(
-        self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context
+        self,
+        mock_set_consume_checkpoint,
+        mock_dsm_set_checkpoint,
+        mock_extract_from_lambda_context,
     ):
         """Test Kinesis event with empty eventSourceARN"""
         kinesis_data = {"message": "test"}
@@ -3085,22 +3107,21 @@ def test_kinesis_event_with_empty_arn(
 
         result = extract_context_from_kinesis_event(event, self.lambda_context)
 
-        mock_dsm_set_checkpoint.assert_not_called()
+        mock_dsm_set_checkpoint.assert_called_once_with(None, "kinesis", "")
+        mock_set_consume_checkpoint.assert_not_called()
         mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context)
         self.assertEqual(result, mock_context)
 
     @patch("datadog_lambda.tracing.extract_context_from_lambda_context")
-    @patch("datadog_lambda.tracing._dsm_set_checkpoint")
+    @patch("ddtrace.data_streams.set_consume_checkpoint")
     def test_kinesis_event_with_exception_accessing_first_record(
-        self, mock_dsm_set_checkpoint, mock_extract_from_lambda_context
+        self, mock_set_consume_checkpoint, mock_extract_from_lambda_context
     ):
         event = {"Records": None}
 
         mock_context = Context(trace_id=123, span_id=456)
         mock_extract_from_lambda_context.return_value = mock_context
 
         result = extract_context_from_kinesis_event(event, self.lambda_context)
-
-        mock_dsm_set_checkpoint.assert_not_called()
-        mock_extract_from_lambda_context.assert_called_once_with(self.lambda_context)
+        mock_set_consume_checkpoint.assert_not_called()
         self.assertEqual(result, mock_context)