Yasuke/moto
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix AttributeNames for sqs.receive_message (#3736)

Browse Source 
* Fix AttributeNames for sqs.receive_message

* Fix Lambda issue

* Change to parametrized tests

* Simplify attribute logic
This commit is contained in:
Anton Grübel 2021-03-05 11:42:07 +01:00 committed by GitHub
parent ac0b4bbcf7
commit 6da4905da9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 383 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
moto/core/utils.py
View File

@ -62,6 +62,11 @@ def pascal_to_camelcase(argument):
return argument[0].lower() + argument[1:] return argument[0].lower() + argument[1:]
def camelcase_to_pascal(argument):
"""Converts a camelCase param to the PascalCase equivalent"""
return argument[0].upper() + argument[1:]
def method_names_from_class(clazz): def method_names_from_class(clazz):
# On Python 2, methods are different from functions, and the `inspect` # On Python 2, methods are different from functions, and the `inspect`
# predicates distinguish between them. On Python 3, methods are just # predicates distinguish between them. On Python 3, methods are just

7
moto/sqs/models.py
View File

@ -3,7 +3,10 @@ from __future__ import unicode_literals
import base64 import base64
import hashlib import hashlib
import json import json
import random
import re import re
import string
import six import six
import struct import struct
from copy import deepcopy from copy import deepcopy
@ -78,6 +81,7 @@ class Message(BaseModel):
self.approximate_receive_count = 0 self.approximate_receive_count = 0
self.deduplication_id = None self.deduplication_id = None
self.group_id = None self.group_id = None
self.sequence_number = None
self.visible_at = 0 self.visible_at = 0
self.delayed_until = 0 self.delayed_until = 0
@ -697,6 +701,9 @@ class SQSBackend(BaseBackend):
# Attributes, but not *message* attributes # Attributes, but not *message* attributes
if deduplication_id is not None: if deduplication_id is not None:
message.deduplication_id = deduplication_id message.deduplication_id = deduplication_id
message.sequence_number = "".join(
random.choice(string.digits) for _ in range(20)
)
if group_id is not None: if group_id is not None:
message.group_id = group_id message.group_id = group_id

47
moto/sqs/responses.py
View File

@ -3,7 +3,12 @@ from __future__ import unicode_literals
import re import re
from moto.core.responses import BaseResponse from moto.core.responses import BaseResponse
from moto.core.utils import amz_crc32, amzn_request_id from moto.core.utils import (
amz_crc32,
amzn_request_id,
underscores_to_camelcase,
camelcase_to_pascal,
)
from six.moves.urllib.parse import urlparse from six.moves.urllib.parse import urlparse
from .exceptions import ( from .exceptions import (
@ -354,7 +359,9 @@ class SQSResponse(BaseResponse):
queue_name = self._get_queue_name() queue_name = self._get_queue_name()
message_attributes = self._get_multi_param("message_attributes") message_attributes = self._get_multi_param("message_attributes")
if not message_attributes: if not message_attributes:
message_attributes = extract_input_message_attributes(self.querystring,) message_attributes = extract_input_message_attributes(self.querystring)
attribute_names = self._get_multi_param("AttributeName")
queue = self.sqs_backend.get_queue(queue_name) queue = self.sqs_backend.get_queue(queue_name)
@ -396,8 +403,24 @@ class SQSResponse(BaseResponse):
messages = self.sqs_backend.receive_messages( messages = self.sqs_backend.receive_messages(
queue_name, message_count, wait_time, visibility_timeout, message_attributes queue_name, message_count, wait_time, visibility_timeout, message_attributes
) )
attributes = {
"approximate_first_receive_timestamp": False,
"approximate_receive_count": False,
"message_deduplication_id": False,
"message_group_id": False,
"sender_id": False,
"sent_timestamp": False,
"sequence_number": False,
}
for attribute in attributes:
pascalcase_name = camelcase_to_pascal(underscores_to_camelcase(attribute))
if any(x in ["All", pascalcase_name] for x in attribute_names):
attributes[attribute] = True
template = self.response_template(RECEIVE_MESSAGE_RESPONSE) template = self.response_template(RECEIVE_MESSAGE_RESPONSE)
return template.render(messages=messages) return template.render(messages=messages, attributes=attributes)
def list_dead_letter_source_queues(self): def list_dead_letter_source_queues(self):
request_url = urlparse(self.uri) request_url = urlparse(self.uri)
@ -537,34 +560,48 @@ RECEIVE_MESSAGE_RESPONSE = """<ReceiveMessageResponse>
<ReceiptHandle>{{ message.receipt_handle }}</ReceiptHandle> <ReceiptHandle>{{ message.receipt_handle }}</ReceiptHandle>
<MD5OfBody>{{ message.body_md5 }}</MD5OfBody> <MD5OfBody>{{ message.body_md5 }}</MD5OfBody>
<Body>{{ message.body }}</Body> <Body>{{ message.body }}</Body>
{% if attributes.sender_id %}
<Attribute> <Attribute>
<Name>SenderId</Name> <Name>SenderId</Name>
<Value>{{ message.sender_id }}</Value> <Value>{{ message.sender_id }}</Value>
</Attribute> </Attribute>
{% endif %}
{% if attributes.sent_timestamp %}
<Attribute> <Attribute>
<Name>SentTimestamp</Name> <Name>SentTimestamp</Name>
<Value>{{ message.sent_timestamp }}</Value> <Value>{{ message.sent_timestamp }}</Value>
</Attribute> </Attribute>
{% endif %}
{% if attributes.approximate_receive_count %}
<Attribute> <Attribute>
<Name>ApproximateReceiveCount</Name> <Name>ApproximateReceiveCount</Name>
<Value>{{ message.approximate_receive_count }}</Value> <Value>{{ message.approximate_receive_count }}</Value>
</Attribute> </Attribute>
{% endif %}
{% if attributes.approximate_first_receive_timestamp %}
<Attribute> <Attribute>
<Name>ApproximateFirstReceiveTimestamp</Name> <Name>ApproximateFirstReceiveTimestamp</Name>
<Value>{{ message.approximate_first_receive_timestamp }}</Value> <Value>{{ message.approximate_first_receive_timestamp }}</Value>
</Attribute> </Attribute>
{% if message.deduplication_id is not none %} {% endif %}
{% if attributes.message_deduplication_id and message.deduplication_id is not none %}
<Attribute> <Attribute>
<Name>MessageDeduplicationId</Name> <Name>MessageDeduplicationId</Name>
<Value>{{ message.deduplication_id }}</Value> <Value>{{ message.deduplication_id }}</Value>
</Attribute> </Attribute>
{% endif %} {% endif %}
{% if message.group_id is not none %} {% if attributes.message_group_id and message.group_id is not none %}
<Attribute> <Attribute>
<Name>MessageGroupId</Name> <Name>MessageGroupId</Name>
<Value>{{ message.group_id }}</Value> <Value>{{ message.group_id }}</Value>
</Attribute> </Attribute>
{% endif %} {% endif %}
{% if attributes.sequence_number and message.sequence_number is not none %}
<Attribute>
<Name>SequenceNumber</Name>
<Value>{{ message.sequence_number }}</Value>
</Attribute>
{% endif %}
{% if message.message_attributes.items()|count > 0 %} {% if message.message_attributes.items()|count > 0 %}
<MD5OfMessageAttributes>{{- message.attribute_md5 -}}</MD5OfMessageAttributes> <MD5OfMessageAttributes>{{- message.attribute_md5 -}}</MD5OfMessageAttributes>
{% endif %} {% endif %}

49
tests/test_core/test_utils.py
View File

@ -3,6 +3,7 @@ from __future__ import unicode_literals
import copy import copy
import sys import sys
import pytest
import sure # noqa import sure # noqa
from freezegun import freeze_time from freezegun import freeze_time
@ -11,24 +12,46 @@ from moto.core.utils import (
underscores_to_camelcase, underscores_to_camelcase,
unix_time, unix_time,
py2_strip_unicode_keys, py2_strip_unicode_keys,
camelcase_to_pascal,
pascal_to_camelcase,
) )
def test_camelcase_to_underscores(): @pytest.mark.parametrize(
cases = { "input,expected",
"theNewAttribute": "the_new_attribute", [
"attri bute With Space": "attribute_with_space", ("theNewAttribute", "the_new_attribute"),
"FirstLetterCapital": "first_letter_capital", ("attri bute With Space", "attribute_with_space"),
"ListMFADevices": "list_mfa_devices", ("FirstLetterCapital", "first_letter_capital"),
} ("ListMFADevices", "list_mfa_devices"),
for arg, expected in cases.items(): ],
camelcase_to_underscores(arg).should.equal(expected) )
def test_camelcase_to_underscores(input, expected):
camelcase_to_underscores(input).should.equal(expected)
def test_underscores_to_camelcase(): @pytest.mark.parametrize(
cases = {"the_new_attribute": "theNewAttribute"} "input,expected",
for arg, expected in cases.items(): [("the_new_attribute", "theNewAttribute"), ("attribute", "attribute"),],
underscores_to_camelcase(arg).should.equal(expected) )
def test_underscores_to_camelcase(input, expected):
underscores_to_camelcase(input).should.equal(expected)
@pytest.mark.parametrize(
"input,expected",
[("TheNewAttribute", "theNewAttribute"), ("Attribute", "attribute"),],
)
def test_pascal_to_camelcase(input, expected):
pascal_to_camelcase(input).should.equal(expected)
@pytest.mark.parametrize(
"input,expected",
[("theNewAttribute", "TheNewAttribute"), ("attribute", "Attribute"),],
)
def test_camelcase_to_pascal(input, expected):
camelcase_to_pascal(input).should.equal(expected)
@freeze_time("2015-01-01 12:00:00") @freeze_time("2015-01-01 12:00:00")

299
tests/test_sqs/test_sqs.py
View File

@ -442,7 +442,9 @@ def test_send_message_with_message_group_id():
MessageGroupId="group_id_1", MessageGroupId="group_id_1",
) )
messages = queue.receive_messages() messages = queue.receive_messages(
AttributeNames=["MessageDeduplicationId", "MessageGroupId"]
)
messages.should.have.length_of(1) messages.should.have.length_of(1)
message_attributes = messages[0].attributes message_attributes = messages[0].attributes
@ -670,6 +672,9 @@ def test_send_receive_message_without_attributes():
message1.shouldnt.have.key("MD5OfMessageAttributes") message1.shouldnt.have.key("MD5OfMessageAttributes")
message2.shouldnt.have.key("MD5OfMessageAttributes") message2.shouldnt.have.key("MD5OfMessageAttributes")
message1.should_not.have.key("Attributes")
message2.should_not.have.key("Attributes")
@mock_sqs @mock_sqs
def test_send_receive_message_with_attributes(): def test_send_receive_message_with_attributes():
@ -782,9 +787,11 @@ def test_send_receive_message_timestamps():
response = queue.send_message(MessageBody="derp") response = queue.send_message(MessageBody="derp")
assert response["ResponseMetadata"]["RequestId"] assert response["ResponseMetadata"]["RequestId"]
messages = conn.receive_message(QueueUrl=queue.url, MaxNumberOfMessages=1)[ messages = conn.receive_message(
"Messages" QueueUrl=queue.url,
] AttributeNames=["ApproximateFirstReceiveTimestamp", "SentTimestamp"],
MaxNumberOfMessages=1,
)["Messages"]
message = messages[0] message = messages[0]
sent_timestamp = message.get("Attributes").get("SentTimestamp") sent_timestamp = message.get("Attributes").get("SentTimestamp")
@ -796,6 +803,283 @@ def test_send_receive_message_timestamps():
int.when.called_with(approximate_first_receive_timestamp).shouldnt.throw(ValueError) int.when.called_with(approximate_first_receive_timestamp).shouldnt.throw(ValueError)
@mock_sqs
@pytest.mark.parametrize(
"attribute_name,expected",
[
(
"All",
{
"ApproximateFirstReceiveTimestamp": lambda x: x.should_not.be.empty,
"ApproximateReceiveCount": lambda x: x.should.equal("1"),
"MessageDeduplicationId": lambda x: x.should.be.none,
"MessageGroupId": lambda x: x.should.be.none,
"SenderId": lambda x: x.should_not.be.empty,
"SentTimestamp": lambda x: x.should_not.be.empty,
"SequenceNumber": lambda x: x.should.be.none,
},
),
(
"ApproximateFirstReceiveTimestamp",
{
"ApproximateFirstReceiveTimestamp": lambda x: x.should_not.be.empty,
"ApproximateReceiveCount": lambda x: x.should.be.none,
"MessageDeduplicationId": lambda x: x.should.be.none,
"MessageGroupId": lambda x: x.should.be.none,
"SenderId": lambda x: x.should.be.none,
"SentTimestamp": lambda x: x.should.be.none,
"SequenceNumber": lambda x: x.should.be.none,
},
),
(
"ApproximateReceiveCount",
{
"ApproximateFirstReceiveTimestamp": lambda x: x.should.be.none,
"ApproximateReceiveCount": lambda x: x.should.equal("1"),
"MessageDeduplicationId": lambda x: x.should.be.none,
"MessageGroupId": lambda x: x.should.be.none,
"SenderId": lambda x: x.should.be.none,
"SentTimestamp": lambda x: x.should.be.none,
"SequenceNumber": lambda x: x.should.be.none,
},
),
(
"SenderId",
{
"ApproximateFirstReceiveTimestamp": lambda x: x.should.be.none,
"ApproximateReceiveCount": lambda x: x.should.be.none,
"MessageDeduplicationId": lambda x: x.should.be.none,
"MessageGroupId": lambda x: x.should.be.none,
"SenderId": lambda x: x.should_not.be.empty,
"SentTimestamp": lambda x: x.should.be.none,
"SequenceNumber": lambda x: x.should.be.none,
},
),
(
"SentTimestamp",
{
"ApproximateFirstReceiveTimestamp": lambda x: x.should.be.none,
"ApproximateReceiveCount": lambda x: x.should.be.none,
"MessageDeduplicationId": lambda x: x.should.be.none,
"MessageGroupId": lambda x: x.should.be.none,
"SenderId": lambda x: x.should.be.none,
"SentTimestamp": lambda x: x.should_not.be.empty,
"SequenceNumber": lambda x: x.should.be.none,
},
),
],
ids=[
"All",
"ApproximateFirstReceiveTimestamp",
"ApproximateReceiveCount",
"SenderId",
"SentTimestamp",
],
)
def test_send_receive_message_with_attribute_name(attribute_name, expected):
sqs = boto3.resource("sqs", region_name="us-east-1")
client = boto3.client("sqs", region_name="us-east-1")
client.create_queue(QueueName="test-queue")
queue = sqs.Queue("test-queue")
body_one = "this is a test message"
body_two = "this is another test message"
queue.send_message(MessageBody=body_one)
queue.send_message(MessageBody=body_two)
messages = client.receive_message(
QueueUrl=queue.url, AttributeNames=[attribute_name], MaxNumberOfMessages=2
)["Messages"]
message1 = messages[0]
message2 = messages[1]
message1["Body"].should.equal(body_one)
message2["Body"].should.equal(body_two)
message1.shouldnt.have.key("MD5OfMessageAttributes")
message2.shouldnt.have.key("MD5OfMessageAttributes")
expected["ApproximateFirstReceiveTimestamp"](
message1["Attributes"].get("ApproximateFirstReceiveTimestamp")
)
expected["ApproximateReceiveCount"](
message1["Attributes"].get("ApproximateReceiveCount")
)
expected["MessageDeduplicationId"](
message1["Attributes"].get("MessageDeduplicationId")
)
expected["MessageGroupId"](message1["Attributes"].get("MessageGroupId"))
expected["SenderId"](message1["Attributes"].get("SenderId"))
expected["SentTimestamp"](message1["Attributes"].get("SentTimestamp"))
expected["SequenceNumber"](message1["Attributes"].get("SequenceNumber"))
expected["ApproximateFirstReceiveTimestamp"](
message2["Attributes"].get("ApproximateFirstReceiveTimestamp")
)
expected["ApproximateReceiveCount"](
message2["Attributes"].get("ApproximateReceiveCount")
)
expected["MessageDeduplicationId"](
message2["Attributes"].get("MessageDeduplicationId")
)
expected["MessageGroupId"](message2["Attributes"].get("MessageGroupId"))
expected["SenderId"](message2["Attributes"].get("SenderId"))
expected["SentTimestamp"](message2["Attributes"].get("SentTimestamp"))
expected["SequenceNumber"](message2["Attributes"].get("SequenceNumber"))
@mock_sqs
@pytest.mark.parametrize(
"attribute_name,expected",
[
(
"All",
{
"ApproximateFirstReceiveTimestamp": lambda x: x.should_not.be.empty,
"ApproximateReceiveCount": lambda x: x.should.equal("1"),
"MessageDeduplicationId": lambda x: x.should.equal("123"),
"MessageGroupId": lambda x: x.should.equal("456"),
"SenderId": lambda x: x.should_not.be.empty,
"SentTimestamp": lambda x: x.should_not.be.empty,
"SequenceNumber": lambda x: x.should_not.be.empty,
},
),
(
"ApproximateFirstReceiveTimestamp",
{
"ApproximateFirstReceiveTimestamp": lambda x: x.should_not.be.empty,
"ApproximateReceiveCount": lambda x: x.should.be.none,
"MessageDeduplicationId": lambda x: x.should.be.none,
"MessageGroupId": lambda x: x.should.be.none,
"SenderId": lambda x: x.should.be.none,
"SentTimestamp": lambda x: x.should.be.none,
"SequenceNumber": lambda x: x.should.be.none,
},
),
(
"ApproximateReceiveCount",
{
"ApproximateFirstReceiveTimestamp": lambda x: x.should.be.none,
"ApproximateReceiveCount": lambda x: x.should.equal("1"),
"MessageDeduplicationId": lambda x: x.should.be.none,
"MessageGroupId": lambda x: x.should.be.none,
"SenderId": lambda x: x.should.be.none,
"SentTimestamp": lambda x: x.should.be.none,
"SequenceNumber": lambda x: x.should.be.none,
},
),
(
"MessageDeduplicationId",
{
"ApproximateFirstReceiveTimestamp": lambda x: x.should.be.none,
"ApproximateReceiveCount": lambda x: x.should.be.none,
"MessageDeduplicationId": lambda x: x.should.equal("123"),
"MessageGroupId": lambda x: x.should.be.none,
"SenderId": lambda x: x.should.be.none,
"SentTimestamp": lambda x: x.should.be.none,
"SequenceNumber": lambda x: x.should.be.none,
},
),
(
"MessageGroupId",
{
"ApproximateFirstReceiveTimestamp": lambda x: x.should.be.none,
"ApproximateReceiveCount": lambda x: x.should.be.none,
"MessageDeduplicationId": lambda x: x.should.be.none,
"MessageGroupId": lambda x: x.should.equal("456"),
"SenderId": lambda x: x.should.be.none,
"SentTimestamp": lambda x: x.should.be.none,
"SequenceNumber": lambda x: x.should.be.none,
},
),
(
"SenderId",
{
"ApproximateFirstReceiveTimestamp": lambda x: x.should.be.none,
"ApproximateReceiveCount": lambda x: x.should.be.none,
"MessageDeduplicationId": lambda x: x.should.be.none,
"MessageGroupId": lambda x: x.should.be.none,
"SenderId": lambda x: x.should_not.be.empty,
"SentTimestamp": lambda x: x.should.be.none,
"SequenceNumber": lambda x: x.should.be.none,
},
),
(
"SentTimestamp",
{
"ApproximateFirstReceiveTimestamp": lambda x: x.should.be.none,
"ApproximateReceiveCount": lambda x: x.should.be.none,
"MessageDeduplicationId": lambda x: x.should.be.none,
"MessageGroupId": lambda x: x.should.be.none,
"SenderId": lambda x: x.should.be.none,
"SentTimestamp": lambda x: x.should_not.be.empty,
"SequenceNumber": lambda x: x.should.be.none,
},
),
(
"SequenceNumber",
{
"ApproximateFirstReceiveTimestamp": lambda x: x.should.be.none,
"ApproximateReceiveCount": lambda x: x.should.be.none,
"MessageDeduplicationId": lambda x: x.should.be.none,
"MessageGroupId": lambda x: x.should.be.none,
"SenderId": lambda x: x.should.be.none,
"SentTimestamp": lambda x: x.should.be.none,
"SequenceNumber": lambda x: x.should_not.be.empty,
},
),
],
ids=[
"All",
"ApproximateFirstReceiveTimestamp",
"ApproximateReceiveCount",
"MessageDeduplicationId",
"MessageGroupId",
"SenderId",
"SentTimestamp",
"SequenceNumber",
],
)
def test_fifo_send_receive_message_with_attribute_name(attribute_name, expected):
client = boto3.client("sqs", region_name="us-east-1")
queue_url = client.create_queue(
QueueName="test-queue.fifo", Attributes={"FifoQueue": "true"}
)["QueueUrl"]
body = "this is a test message"
client.send_message(
QueueUrl=queue_url,
MessageBody=body,
MessageDeduplicationId="123",
MessageGroupId="456",
)
message = client.receive_message(
QueueUrl=queue_url, AttributeNames=[attribute_name], MaxNumberOfMessages=2
)["Messages"][0]
message["Body"].should.equal(body)
message.should_not.have.key("MD5OfMessageAttributes")
expected["ApproximateFirstReceiveTimestamp"](
message["Attributes"].get("ApproximateFirstReceiveTimestamp")
)
expected["ApproximateReceiveCount"](
message["Attributes"].get("ApproximateReceiveCount")
)
expected["MessageDeduplicationId"](
message["Attributes"].get("MessageDeduplicationId")
)
expected["MessageGroupId"](message["Attributes"].get("MessageGroupId"))
expected["SenderId"](message["Attributes"].get("SenderId"))
expected["SentTimestamp"](message["Attributes"].get("SentTimestamp"))
expected["SequenceNumber"](message["Attributes"].get("SequenceNumber"))
@mock_sqs @mock_sqs
def test_max_number_of_messages_invalid_param(): def test_max_number_of_messages_invalid_param():
sqs = boto3.resource("sqs", region_name="us-east-1") sqs = boto3.resource("sqs", region_name="us-east-1")
@ -1013,7 +1297,7 @@ def test_message_attributes():
queue.count().should.equal(1) queue.count().should.equal(1)
messages = conn.receive_message(queue, number_messages=1) messages = conn.receive_message(queue, number_messages=1, attributes=["All"])
queue.count().should.equal(0) queue.count().should.equal(0)
assert len(messages) == 1 assert len(messages) == 1
@ -1347,6 +1631,7 @@ def test_send_message_batch():
QueueUrl=queue_url, QueueUrl=queue_url,
MaxNumberOfMessages=10, MaxNumberOfMessages=10,
MessageAttributeNames=["attribute_name_1", "attribute_name_2"], MessageAttributeNames=["attribute_name_1", "attribute_name_2"],
AttributeNames=["MessageDeduplicationId", "MessageGroupId"],
) )
response["Messages"][0]["Body"].should.equal("body_1") response["Messages"][0]["Body"].should.equal("body_1")
@ -2098,7 +2383,9 @@ def test_receive_messages_with_message_group_id():
queue.send_message(MessageBody="message-3", MessageGroupId="group") queue.send_message(MessageBody="message-3", MessageGroupId="group")
queue.send_message(MessageBody="separate-message", MessageGroupId="anothergroup") queue.send_message(MessageBody="separate-message", MessageGroupId="anothergroup")
messages = queue.receive_messages(MaxNumberOfMessages=2) messages = queue.receive_messages(
MaxNumberOfMessages=2, AttributeNames=["MessageGroupId"]
)
messages.should.have.length_of(2) messages.should.have.length_of(2)
messages[0].attributes["MessageGroupId"].should.equal("group") messages[0].attributes["MessageGroupId"].should.equal("group")