Skip to content

Commit c591ea4

Browse files
gregfurmangregfurman
committed
[SQS] Add override header tests
1 parent 374b006 commit c591ea4

File tree

4 files changed

+188
-0
lines changed

4 files changed

+188
-0
lines changed

‎tests/aws/services/sqs/test_sqs.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -263,6 +263,19 @@ def test_receive_empty_queue(self, sqs_queue, snapshot, aws_sqs_client):
263263
QueueUrl=queue_url, MaxNumberOfMessages=1, WaitTimeSeconds=1
264264
)
265265
snapshot.match("empty_long_poll_resp", empty_long_poll_resp)
266+
@markers.snapshot.skip_snapshot_verify(paths=["$..Error.Detail"])
267+
def test_send_receive_wait_time_seconds(self, sqs_queue, snapshot, aws_sqs_client):
268+
queue_url = sqs_queue
269+
send_result = aws_sqs_client.send_message(QueueUrl=queue_url, MessageBody="message")
270+
assert send_result["MessageId"]
271+
272+
MAX_WAIT_TIME_SECONDS = 20
273+
with pytest.raises(ClientError) as e:
274+
aws_sqs_client.receive_message(
275+
QueueUrl=queue_url, WaitTimeSeconds=MAX_WAIT_TIME_SECONDS + 1
276+
)
277+
278+
snapshot.match("recieve_message_error", e.value.response)
266279

267280
@markers.aws.validated
268281
def test_receive_message_attributes_timestamp_types(self, sqs_queue, aws_sqs_client):

‎tests/aws/services/sqs/test_sqs.snapshot.json

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3647,6 +3647,41 @@
36473647
"recorded-date": "20-08-2024, 14:14:11",
36483648
"recorded-content": {}
36493649
},
3650+
"tests/aws/services/sqs/test_sqs.py::TestSqsProvider::test_send_receive_wait_time_seconds[sqs]": {
3651+
"recorded-date": "23-01-2025, 12:43:53",
3652+
"recorded-content": {
3653+
"recieve_message_error": {
3654+
"Error": {
3655+
"Code": "InvalidParameterValue",
3656+
"Message": "Value 21 for parameter WaitTimeSeconds is invalid. Reason: Must be >= 0 and <= 20, if provided.",
3657+
"QueryErrorCode": "InvalidParameterValueException",
3658+
"Type": "Sender"
3659+
},
3660+
"ResponseMetadata": {
3661+
"HTTPHeaders": {},
3662+
"HTTPStatusCode": 400
3663+
3664+
}
3665+
}
3666+
}
3667+
},
3668+
"tests/aws/services/sqs/test_sqs.py::TestSqsProvider::test_send_receive_wait_time_seconds[sqs_query]": {
3669+
"recorded-date": "23-01-2025, 12:43:54",
3670+
"recorded-content": {
3671+
"recieve_message_error": {
3672+
"Error": {
3673+
"Code": "InvalidParameterValue",
3674+
"Detail": null,
3675+
"Message": "Value 21 for parameter WaitTimeSeconds is invalid. Reason: Must be >= 0 and <= 20, if provided.",
3676+
"Type": "Sender"
3677+
},
3678+
"ResponseMetadata": {
3679+
"HTTPHeaders": {},
3680+
"HTTPStatusCode": 400
3681+
}
3682+
}
3683+
}
3684+
},
36503685
"tests/aws/services/sqs/test_sqs.py::TestSqsProvider::test_receive_empty_queue[sqs]": {
36513686
"recorded-date": "30-01-2025, 22:32:45",
36523687
"recorded-content": {

‎tests/aws/services/sqs/test_sqs.validation.json

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -329,6 +329,12 @@
329329
"tests/aws/services/sqs/test_sqs.py::TestSqsProvider::test_send_receive_message_multiple_queues": {
330330
"last_validated_date": "2024-04-30T13:40:05+00:00"
331331
},
332+
"tests/aws/services/sqs/test_sqs.py::TestSqsProvider::test_send_receive_wait_time_seconds[sqs]": {
333+
"last_validated_date": "2025-01-23T12:43:53+00:00"
334+
},
335+
"tests/aws/services/sqs/test_sqs.py::TestSqsProvider::test_send_receive_wait_time_seconds[sqs_query]": {
336+
"last_validated_date": "2025-01-23T12:43:54+00:00"
337+
},
332338
"tests/aws/services/sqs/test_sqs.py::TestSqsProvider::test_set_empty_redrive_policy[sqs]": {
333339
"last_validated_date": "2024-08-20T14:14:08+00:00"
334340
},
@@ -401,6 +407,12 @@
401407
"tests/aws/services/sqs/test_sqs.py::TestSqsProvider::test_too_many_entries_in_batch_request[sqs_query]": {
402408
"last_validated_date": "2024-04-30T13:33:40+00:00"
403409
},
410+
"tests/aws/services/sqs/test_sqs.py::TestSqsProvider::test_wait_time_seconds_waits_correctly[sqs]": {
411+
"last_validated_date": "2025-01-23T13:57:19+00:00"
412+
},
413+
"tests/aws/services/sqs/test_sqs.py::TestSqsProvider::test_wait_time_seconds_waits_correctly[sqs_query]": {
414+
"last_validated_date": "2025-01-23T13:57:30+00:00"
415+
},
404416
"tests/aws/services/sqs/test_sqs.py::TestSqsQueryApi::test_send_message_via_queue_url_with_json_protocol": {
405417
"last_validated_date": "2024-04-30T13:35:11+00:00"
406418
}

‎tests/aws/services/sqs/test_sqs_backdoor.py

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,6 @@
1+
import time
2+
from threading import Timer
3+
14
import pytest
25
import requests
36
import xmltodict
@@ -361,3 +364,128 @@ def test_list_messages_with_queue_url_in_path(
361364
assert response.status_code == 400
362365
doc = response.json()
363366
assert doc["ErrorResponse"]["Error"]["Code"] == "AWS.SimpleQueueService.NonExistentQueue"
367+
368+
369+
class TestSqsOverrideHeaders:
370+
@markers.aws.only_localstack
371+
def test_receive_message_override_max_number_of_messages(
372+
self, sqs_create_queue, aws_client_factory
373+
):
374+
# Create standalone boto3 client since registering hooks to the session-wide
375+
# aws_client (from the fixture) will have side-effects.
376+
aws_client = aws_client_factory().sqs
377+
378+
override_max_number_of_messages = 20
379+
380+
from localstack.services.sqs.constants import HEADER_LOCALSTACK_SQS_OVERRIDE_MESSAGE_COUNT
381+
from localstack.services.sqs.provider import MAX_NUMBER_OF_MESSAGES
382+
383+
queue_url = sqs_create_queue()
384+
385+
for i in range(override_max_number_of_messages):
386+
aws_client.send_message(QueueUrl=queue_url, MessageBody=f"message-{i}")
387+
388+
with pytest.raises(ClientError):
389+
aws_client.receive_message(
390+
QueueUrl=queue_url,
391+
VisibilityTimeout=0,
392+
MaxNumberOfMessages=override_max_number_of_messages,
393+
AttributeNames=["All"],
394+
)
395+
396+
def _handle_receive_message_override(params, context, **kwargs):
397+
if not (requested_count := params.get("MaxNumberOfMessages")):
398+
return
399+
context[HEADER_LOCALSTACK_SQS_OVERRIDE_MESSAGE_COUNT] = str(requested_count)
400+
params["MaxNumberOfMessages"] = min(MAX_NUMBER_OF_MESSAGES, requested_count)
401+
402+
def _handler_inject_header(params, context, **kwargs):
403+
if override_message_count := context.pop(
404+
HEADER_LOCALSTACK_SQS_OVERRIDE_MESSAGE_COUNT, None
405+
):
406+
params["headers"][HEADER_LOCALSTACK_SQS_OVERRIDE_MESSAGE_COUNT] = (
407+
override_message_count
408+
)
409+
410+
aws_client.meta.events.register(
411+
"provide-client-params.sqs.ReceiveMessage", _handle_receive_message_override
412+
)
413+
414+
aws_client.meta.events.register("before-call.sqs.ReceiveMessage", _handler_inject_header)
415+
416+
response = aws_client.receive_message(
417+
QueueUrl=queue_url,
418+
VisibilityTimeout=30,
419+
MaxNumberOfMessages=override_max_number_of_messages,
420+
AttributeNames=["All"],
421+
)
422+
423+
messages = response.get("Messages", [])
424+
assert len(messages) == 20
425+
426+
@markers.aws.only_localstack
427+
def test_receive_message_override_message_wait_time_seconds(
428+
self, sqs_create_queue, aws_client_factory
429+
):
430+
aws_client = aws_client_factory().sqs
431+
432+
override_message_wait_time_seconds = 30
433+
434+
from localstack.services.sqs.constants import (
435+
HEADER_LOCALSTACK_SQS_OVERRIDE_WAIT_TIME_SECONDS,
436+
)
437+
from localstack.services.sqs.provider import MAX_NUMBER_OF_MESSAGES
438+
439+
queue_url = sqs_create_queue()
440+
441+
with pytest.raises(ClientError):
442+
aws_client.receive_message(
443+
QueueUrl=queue_url,
444+
VisibilityTimeout=0,
445+
MaxNumberOfMessages=MAX_NUMBER_OF_MESSAGES,
446+
WaitTimeSeconds=override_message_wait_time_seconds,
447+
AttributeNames=["All"],
448+
)
449+
450+
def _handle_receive_message_override(params, context, **kwargs):
451+
if not (requested_wait_time := params.get("WaitTimeSeconds")):
452+
return
453+
context[HEADER_LOCALSTACK_SQS_OVERRIDE_WAIT_TIME_SECONDS] = str(requested_wait_time)
454+
params["WaitTimeSeconds"] = min(20, requested_wait_time)
455+
456+
def _handler_inject_header(params, context, **kwargs):
457+
if override_wait_time := context.pop(
458+
HEADER_LOCALSTACK_SQS_OVERRIDE_WAIT_TIME_SECONDS, None
459+
):
460+
params["headers"][HEADER_LOCALSTACK_SQS_OVERRIDE_WAIT_TIME_SECONDS] = (
461+
override_wait_time
462+
)
463+
464+
aws_client.meta.events.register(
465+
"provide-client-params.sqs.ReceiveMessage", _handle_receive_message_override
466+
)
467+
468+
aws_client.meta.events.register("before-call.sqs.ReceiveMessage", _handler_inject_header)
469+
470+
def _send_message():
471+
aws_client.send_message(QueueUrl=queue_url, MessageBody=f"message-{short_uid()}")
472+
473+
# Populate with 9 messages (1 below the MaxNumberOfMessages threshold).
474+
# This should cause long-polling to exit since MaxNumberOfMessages is met.
475+
for _ in range(9):
476+
_send_message()
477+
478+
Timer(25, _send_message).start() # send message asynchronously after 1 second
479+
480+
start_t = time.time()
481+
response = aws_client.receive_message(
482+
QueueUrl=queue_url,
483+
VisibilityTimeout=30,
484+
MaxNumberOfMessages=MAX_NUMBER_OF_MESSAGES,
485+
WaitTimeSeconds=override_message_wait_time_seconds,
486+
AttributeNames=["All"],
487+
)
488+
assert time.time() - start_t >= 25
489+
490+
messages = response.get("Messages", [])
491+
assert len(messages) == 10

0 commit comments

Comments
 (0)