Fix AttributeNames for sqs.receive_message (#3736)

* Fix AttributeNames for sqs.receive_message

* Fix Lambda issue

* Change to parametrized tests

* Simplify attribute logic

tests/test_sqs/test_sqs.py

@@ -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")