Fix:SQS Receive Message (MessageAttributes) in response (#3303)

* Fix:SQS Receive Message (MessageAttributes) in response

* Fixed tests

Co-authored-by: usmankb <usman@krazybee.com>

moto/sqs/models.py

@@ -526,6 +526,14 @@ class Queue(CloudFormationModel):
             }
 
 
+def _filter_message_attributes(message, input_message_attributes):
+    filtered_message_attributes = {}
+    for key, value in message.message_attributes.items():
+        if key in input_message_attributes:
+            filtered_message_attributes[key] = value
+    message.message_attributes = filtered_message_attributes
+
+
 class SQSBackend(BaseBackend):
     def __init__(self, region_name):
         self.region_name = region_name
@@ -718,7 +726,12 @@ class SQSBackend(BaseBackend):
         return None
 
     def receive_messages(
-        self, queue_name, count, wait_seconds_timeout, visibility_timeout
+        self,
+        queue_name,
+        count,
+        wait_seconds_timeout,
+        visibility_timeout,
+        message_attribute_names=None,
     ):
         """
         Attempt to retrieve visible messages from a queue.
@@ -734,6 +747,8 @@ class SQSBackend(BaseBackend):
         :param int wait_seconds_timeout:  The duration (in seconds) for which the call waits for a message to arrive in
          the queue before returning. If a message is available, the call returns sooner than WaitTimeSeconds
         """
+        if message_attribute_names is None:
+            message_attribute_names = []
         queue = self.get_queue(queue_name)
         result = []
         previous_result_count = len(result)
@@ -775,6 +790,7 @@ class SQSBackend(BaseBackend):
                     continue
 
                 message.mark_received(visibility_timeout=visibility_timeout)
+                _filter_message_attributes(message, message_attribute_names)
                 result.append(message)
                 if len(result) >= count:
                     break

moto/sqs/responses.py

@@ -13,7 +13,7 @@ from .exceptions import (
     ReceiptHandleIsInvalid,
 )
 from .models import sqs_backends
-from .utils import parse_message_attributes
+from .utils import parse_message_attributes, extract_input_message_attributes
 
 MAXIMUM_VISIBILTY_TIMEOUT = 43200
 MAXIMUM_MESSAGE_LENGTH = 262144  # 256 KiB
@@ -352,6 +352,9 @@ class SQSResponse(BaseResponse):
 
     def receive_message(self):
         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,)
 
         queue = self.sqs_backend.get_queue(queue_name)
 
@@ -391,7 +394,7 @@ class SQSResponse(BaseResponse):
             return ERROR_MAX_VISIBILITY_TIMEOUT_RESPONSE, dict(status=400)
 
         messages = self.sqs_backend.receive_messages(
-            queue_name, message_count, wait_time, visibility_timeout
+            queue_name, message_count, wait_time, visibility_timeout, message_attributes
         )
         template = self.response_template(RECEIVE_MESSAGE_RESPONSE)
         return template.render(messages=messages)

moto/sqs/utils.py

@@ -11,6 +11,21 @@ def generate_receipt_handle():
     return "".join(random.choice(string.ascii_lowercase) for x in range(length))
 
 
+def extract_input_message_attributes(querystring):
+    message_attributes = []
+    index = 1
+    while True:
+        # Loop through looking for message attributes
+        name_key = "MessageAttributeName.{0}".format(index)
+        name = querystring.get(name_key)
+        if not name:
+            # Found all attributes
+            break
+        message_attributes.append(name[0])
+        index = index + 1
+    return message_attributes
+
+
 def parse_message_attributes(querystring, base="", value_namespace="Value."):
     message_attributes = {}
     index = 1