[ESM] Fix flaky SQS ReportBatchItemFailures test with proper visiblity timeouts (#12323)

gregfurman · web-flow · commit abab2268ab98 · 2025-03-03T15:24:02.000+02:00
diff --git a/tests/aws/services/lambda_/event_source_mapping/test_lambda_integration_sqs.py b/tests/aws/services/lambda_/event_source_mapping/test_lambda_integration_sqs.py
@@ -389,7 +389,6 @@ def receive_dlq():
     snapshot.match("messages", messages)
 
 
-# TODO: flaky against AWS
 @markers.aws.validated
 def test_report_batch_item_failures(
     create_lambda_function,
@@ -412,19 +411,27 @@ def test_report_batch_item_failures(
         "get_destination_queue_url", aws_client.sqs.get_queue_url(QueueName=destination_queue_name)
     )
 
-    # timeout in seconds, used for both the lambda and the queue visibility timeout.
-    # increase to 10 if testing against AWS fails.
-    retry_timeout = 8
+    # If an SQS queue is not receiving a lot of traffic, Lambda can take up to 20s between invocations.
+    # See AWS docs https://docs.aws.amazon.com/lambda/latest/dg/with-sqs.html.
+    retry_timeout = 6
+    visibility_timeout = 8
     retries = 2
 
+    # AWS recommends a visibility timeout should be x6 a Lambda's retry timeout. To ensure a short test
+    # runtime, we just want to ensure messages are re-queued a couple of seconda after any potential timeouts.
+    # See https://docs.aws.amazon.com/lambda/latest/dg/services-sqs-configure.html#events-sqs-queueconfig
+    assert visibility_timeout > retry_timeout, (
+        "A lambda needs to finish processing prior to re-queuing invisible messages"
+    )
+
     # set up lambda function
     function_name = f"failing-lambda-{short_uid()}"
     create_lambda_function(
         func_name=function_name,
         handler_file=LAMBDA_SQS_BATCH_ITEM_FAILURE_FILE,
         runtime=Runtime.python3_12,
         role=lambda_su_role,
-        timeout=retry_timeout,  # timeout needs to be <= than visibility timeout
+        timeout=retry_timeout,
         envvars={"DESTINATION_QUEUE_URL": destination_url},
     )
 
@@ -441,7 +448,7 @@ def test_report_batch_item_failures(
         Attributes={
             "FifoQueue": "true",
             # the visibility timeout is implicitly also the time between retries
-            "VisibilityTimeout": str(retry_timeout),
+            "VisibilityTimeout": str(visibility_timeout),
             "RedrivePolicy": json.dumps(
                 {"deadLetterTargetArn": event_dlq_arn, "maxReceiveCount": retries}
             ),
@@ -521,8 +528,14 @@ def test_report_batch_item_failures(
     assert "Messages" not in dlq_messages or dlq_messages["Messages"] == []
 
     # now wait for the second invocation result which is expected to have processed message 2 and 3
+    # Since we are re-queuing twice, with a visiblity timeout of 8s, this should instead be waiting for 20s => 8s x 2 retries (+ 4s margin).
+    # See AWS docs: https://docs.aws.amazon.com/AWSSimpleQueueService/latest/APIReference/API_ReceiveMessage.html#API_ReceiveMessage_RequestSyntax
+    second_timeout_with_margin = (visibility_timeout * 2) + 4
+    assert second_timeout_with_margin <= 20, (
+        "An SQS ReceiveMessage operation cannot wait for more than 20s"
+    )
     second_invocation = aws_client.sqs.receive_message(
-        QueueUrl=destination_url, WaitTimeSeconds=retry_timeout + 2, MaxNumberOfMessages=1
+        QueueUrl=destination_url, WaitTimeSeconds=second_timeout_with_margin, MaxNumberOfMessages=1
     )
     assert "Messages" in second_invocation
     # hack to make snapshot work
diff --git a/tests/aws/services/lambda_/event_source_mapping/test_lambda_integration_sqs.snapshot.json b/tests/aws/services/lambda_/event_source_mapping/test_lambda_integration_sqs.snapshot.json
@@ -240,7 +240,7 @@
     }
   },
   "tests/aws/services/lambda_/event_source_mapping/test_lambda_integration_sqs.py::test_report_batch_item_failures": {
-    "recorded-date": "12-10-2024, 13:34:15",
+    "recorded-date": "03-03-2025, 11:31:17",
     "recorded-content": {
       "get_destination_queue_url": {
         "QueueUrl": "<queue-url:1>",
diff --git a/tests/aws/services/lambda_/event_source_mapping/test_lambda_integration_sqs.validation.json b/tests/aws/services/lambda_/event_source_mapping/test_lambda_integration_sqs.validation.json
@@ -117,7 +117,7 @@
     "last_validated_date": "2024-10-12T13:33:27+00:00"
   },
   "tests/aws/services/lambda_/event_source_mapping/test_lambda_integration_sqs.py::test_report_batch_item_failures": {
-    "last_validated_date": "2024-10-12T13:34:12+00:00"
+    "last_validated_date": "2025-03-03T11:31:14+00:00"
   },
   "tests/aws/services/lambda_/event_source_mapping/test_lambda_integration_sqs.py::test_report_batch_item_failures_empty_json_batch_succeeds": {
     "last_validated_date": "2024-10-12T13:35:40+00:00"

Original file line number	Diff line number	Diff line change
`@@ -240,7 +240,7 @@`
`240`	`240`	`}`
`241`	`241`	`},`
`242`	`242`	`"tests/aws/services/lambda_/event_source_mapping/test_lambda_integration_sqs.py::test_report_batch_item_failures": {`
`243`		`- "recorded-date": "12-10-2024, 13:34:15",`
	`243`	`+ "recorded-date": "03-03-2025, 11:31:17",`
`244`	`244`	`"recorded-content": {`
`245`	`245`	`"get_destination_queue_url": {`
`246`	`246`	`"QueueUrl": "<queue-url:1>",`