add test for try catch on checkpoint exception

Author: michael-zhao459

tests/test_dsm.py

@@ -141,6 +141,35 @@ def test_sqs_multiple_records_process_each_record(self):
             self.assertEqual(pathway_ctx, expected_contexts[i])
 
     def test_set_context_exception_handled(self):
+        """Test that exceptions in set_consume_checkpoint are properly handled"""
+        self.mock_get_dsm_context_from_sqs_lambda.return_value = {
+            "dd-pathway-ctx": "test-context"
+        }
+
+        self.mock_set_consume_checkpoint.side_effect = Exception("Checkpoint error")
+
+        event = {
+            "Records": [
+                {
+                    "eventSourceARN": "arn:aws:sqs:us-east-1:123456789012:my-queue",
+                    "body": "Test message",
+                    "messageAttributes": {
+                        "_datadog": {
+                            "stringValue": '{"dd-pathway-ctx": "test-context"}',
+                            "dataType": "String",
+                        }
+                    },
+                }
+            ]
+        }
+
+        _dsm_set_sqs_context(event)
+
+        self.mock_logger.error.assert_called_once_with(
+            "Unable to set dsm context: Checkpoint error"
+        )
+
+    def test_get_context_exception_handled(self):
         """Test that exceptions in _get_dsm_context_from_sqs_lambda are properly handled"""
         self.mock_get_dsm_context_from_sqs_lambda.side_effect = Exception(
             "JSON parsing error"
@@ -153,7 +182,7 @@ def test_set_context_exception_handled(self):
                     "body": "Test message",
                     "messageAttributes": {
                         "_datadog": {
-                            "stringValue": "invalid json",
+                            "stringValue": "invalid json{",
                             "dataType": "String",
                         }
                     },
@@ -167,8 +196,6 @@ def test_set_context_exception_handled(self):
             "Unable to set dsm context: JSON parsing error"
         )
 
-        self.mock_set_consume_checkpoint.assert_not_called()
-
 
 class TestGetDSMContextFromSQS(unittest.TestCase):
     def test_sqs_to_lambda_string_value_format(self):