Merge pull request #1574 from wblackconv/add-msg-attributes

Include SNS message attributes with message body when delivering to SQS.

moto/sns/models.py

@@ -94,7 +94,7 @@ class Subscription(BaseModel):
         if self.protocol == 'sqs':
             queue_name = self.endpoint.split(":")[-1]
             region = self.endpoint.split(":")[3]
-            enveloped_message = json.dumps(self.get_post_data(message, message_id, subject), sort_keys=True, indent=2, separators=(',', ': '))
+            enveloped_message = json.dumps(self.get_post_data(message, message_id, subject, message_attributes=message_attributes), sort_keys=True, indent=2, separators=(',', ': '))
             sqs_backends[region].send_message(queue_name, enveloped_message)
         elif self.protocol in ['http', 'https']:
             post_data = self.get_post_data(message, message_id, subject)
@@ -132,15 +132,16 @@ class Subscription(BaseModel):
             for rule in rules:
                 if isinstance(rule, six.string_types):
                     # only string value matching is supported
-                    if message_attributes[field] == rule:
+                    if message_attributes[field]['Value'] == rule:
                         return True
             return False
 
         return all(_field_match(field, rules, message_attributes)
                    for field, rules in six.iteritems(self._filter_policy))
 
-    def get_post_data(self, message, message_id, subject):
-        return {
+    def get_post_data(
+            self, message, message_id, subject, message_attributes=None):
+        post_data = {
             "Type": "Notification",
             "MessageId": message_id,
             "TopicArn": self.topic.arn,
@@ -152,6 +153,9 @@ class Subscription(BaseModel):
             "SigningCertURL": "https://sns.us-east-1.amazonaws.com/SimpleNotificationService-f3ecfb7224c7233fe7bb5f59f96de52f.pem",
             "UnsubscribeURL": "https://sns.us-east-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:us-east-1:123456789012:some-topic:2bcfbf39-05c3-41de-beaa-fcfcc21c8f55"
         }
+        if message_attributes:
+            post_data["MessageAttributes"] = message_attributes
+        return post_data
 
 
 class PlatformApplication(BaseModel):

moto/sns/responses.py

@@ -6,7 +6,7 @@ from collections import defaultdict
 from moto.core.responses import BaseResponse
 from moto.core.utils import camelcase_to_underscores
 from .models import sns_backends
-from .exceptions import SNSNotFoundError
+from .exceptions import SNSNotFoundError, InvalidParameterValue
 from .utils import is_e164
 
 
@@ -30,6 +30,48 @@ class SNSResponse(BaseResponse):
             in attributes
         )
 
+    def _parse_message_attributes(self, prefix='', value_namespace='Value.'):
+        message_attributes = self._get_object_map(
+            'MessageAttributes.entry',
+            name='Name',
+            value='Value'
+        )
+        # SNS converts some key names before forwarding messages
+        # DataType -> Type, StringValue -> Value, BinaryValue -> Value
+        transformed_message_attributes = {}
+        for name, value in message_attributes.items():
+            # validation
+            data_type = value['DataType']
+            if not data_type:
+                raise InvalidParameterValue(
+                    "The message attribute '{0}' must contain non-empty "
+                    "message attribute value.".format(name))
+
+            data_type_parts = data_type.split('.')
+            if (len(data_type_parts) > 2 or
+                    data_type_parts[0] not in ['String', 'Binary', 'Number']):
+                raise InvalidParameterValue(
+                    "The message attribute '{0}' has an invalid message "
+                    "attribute type, the set of supported type prefixes is "
+                    "Binary, Number, and String.".format(name))
+
+            if 'StringValue' in value:
+                value = value['StringValue']
+            elif 'BinaryValue' in 'Value':
+                value = value['BinaryValue']
+            if not value:
+                raise InvalidParameterValue(
+                    "The message attribute '{0}' must contain non-empty "
+                    "message attribute value for message attribute "
+                    "type '{1}'.".format(name, data_type[0]))
+
+            # transformation
+            transformed_message_attributes[name] = {
+                'Type': data_type, 'Value': value
+            }
+
+        return transformed_message_attributes
+
     def create_topic(self):
         name = self._get_param('Name')
         topic = self.backend.create_topic(name)
@@ -241,9 +283,10 @@ class SNSResponse(BaseResponse):
         phone_number = self._get_param('PhoneNumber')
         subject = self._get_param('Subject')
 
-        message_attributes = self._get_map_prefix('MessageAttributes.entry',
-                                                  key_end='Name',
-                                                  value_end='Value')
+        try:
+            message_attributes = self._parse_message_attributes()
+        except InvalidParameterValue as e:
+            return self._error(e.description), dict(status=e.code)
 
         if phone_number is not None:
             # Check phone is correct syntax (e164)