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:]
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):
# On Python 2, methods are different from functions, and the `inspect`
# 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 hashlib
import json
import random
import re
import string
import six
import struct
from copy import deepcopy
@ -78,6 +81,7 @@ class Message(BaseModel):
self.approximate_receive_count = 0
self.deduplication_id = None
self.group_id = None
self.sequence_number = None
self.visible_at = 0
self.delayed_until = 0
@ -697,6 +701,9 @@ class SQSBackend(BaseBackend):
# Attributes, but not *message* attributes
if deduplication_id is not None:
message.deduplication_id = deduplication_id
message.sequence_number = "".join(
random.choice(string.digits) for _ in range(20)
)
if group_id is not None:
message.group_id = group_id

47
moto/sqs/responses.py
View File

@ -3,7 +3,12 @@ from __future__ import unicode_literals
import re
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 .exceptions import (
@ -354,7 +359,9 @@ class SQSResponse(BaseResponse):
queue_name = self._get_queue_name()
message_attributes = self._get_multi_param("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)
@ -396,8 +403,24 @@ class SQSResponse(BaseResponse):
messages = self.sqs_backend.receive_messages(
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)
return template.render(messages=messages)
return template.render(messages=messages, attributes=attributes)
def list_dead_letter_source_queues(self):
request_url = urlparse(self.uri)
@ -537,34 +560,48 @@ RECEIVE_MESSAGE_RESPONSE = """<ReceiveMessageResponse>
<ReceiptHandle>{{ message.receipt_handle }}</ReceiptHandle>
<MD5OfBody>{{ message.body_md5 }}</MD5OfBody>
<Body>{{ message.body }}</Body>
{% if attributes.sender_id %}
<Attribute>
<Name>SenderId</Name>
<Value>{{ message.sender_id }}</Value>
</Attribute>
{% endif %}
{% if attributes.sent_timestamp %}
<Attribute>
<Name>SentTimestamp</Name>
<Value>{{ message.sent_timestamp }}</Value>
</Attribute>
{% endif %}
{% if attributes.approximate_receive_count %}
<Attribute>
<Name>ApproximateReceiveCount</Name>
<Value>{{ message.approximate_receive_count }}</Value>
</Attribute>
{% endif %}
{% if attributes.approximate_first_receive_timestamp %}
<Attribute>
<Name>ApproximateFirstReceiveTimestamp</Name>
<Value>{{ message.approximate_first_receive_timestamp }}</Value>
</Attribute>
{% if message.deduplication_id is not none %}
{% endif %}
{% if attributes.message_deduplication_id and message.deduplication_id is not none %}
<Attribute>
<Name>MessageDeduplicationId</Name>
<Value>{{ message.deduplication_id }}</Value>
</Attribute>
{% endif %}
{% if message.group_id is not none %}
{% if attributes.message_group_id and message.group_id is not none %}
<Attribute>
<Name>MessageGroupId</Name>
<Value>{{ message.group_id }}</Value>
</Attribute>
{% 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 %}
<MD5OfMessageAttributes>{{- message.attribute_md5 -}}</MD5OfMessageAttributes>
{% endif %}

49
tests/test_core/test_utils.py
View File

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

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",
)
messages = queue.receive_messages()
messages = queue.receive_messages(
AttributeNames=["MessageDeduplicationId", "MessageGroupId"]
)
messages.should.have.length_of(1)
message_attributes = messages[0].attributes
@ -670,6 +672,9 @@ def test_send_receive_message_without_attributes():
message1.shouldnt.have.key("MD5OfMessageAttributes")
message2.shouldnt.have.key("MD5OfMessageAttributes")
message1.should_not.have.key("Attributes")
message2.should_not.have.key("Attributes")
@mock_sqs
def test_send_receive_message_with_attributes():
@ -782,9 +787,11 @@ def test_send_receive_message_timestamps():
response = queue.send_message(MessageBody="derp")
assert response["ResponseMetadata"]["RequestId"]
messages = conn.receive_message(QueueUrl=queue.url, MaxNumberOfMessages=1)[
"Messages"
]
messages = conn.receive_message(
QueueUrl=queue.url,
AttributeNames=["ApproximateFirstReceiveTimestamp", "SentTimestamp"],
MaxNumberOfMessages=1,
)["Messages"]
message = messages[0]
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)
@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
def test_max_number_of_messages_invalid_param():
sqs = boto3.resource("sqs", region_name="us-east-1")
@ -1013,7 +1297,7 @@ def test_message_attributes():
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)
assert len(messages) == 1
@ -1347,6 +1631,7 @@ def test_send_message_batch():
QueueUrl=queue_url,
MaxNumberOfMessages=10,
MessageAttributeNames=["attribute_name_1", "attribute_name_2"],
AttributeNames=["MessageDeduplicationId", "MessageGroupId"],
)
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="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[0].attributes["MessageGroupId"].should.equal("group")