Merge 5ef414e into bbc48c3

localstack · Oct 23, 2018 · 93fe53d · 93fe53d
2 parents bbc48c3 + 5ef414e
commit 93fe53d
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 1 deletion.
diff --git a/localstack/services/sns/sns_listener.py b/localstack/services/sns/sns_listener.py
@@ -82,7 +82,8 @@ def forward_request(self, method, path, data, headers):
                             try:
                                 sqs_client.send_message(
                                     QueueUrl=queue_url,
-                                    MessageBody=create_sns_message_body(subscriber, req_data)
+                                    MessageBody=create_sns_message_body(subscriber, req_data),
+                                    MessageAttributes=create_sqs_message_attributes(subscriber, message_attributes)
                                 )
                             except Exception as exc:
                                 return make_error(message=str(exc), code=400)
@@ -226,6 +227,23 @@ def create_sns_message_body(subscriber, req_data):
     return json.dumps(data)
 
 
+def create_sqs_message_attributes(subscriber, attributes):
+    if subscriber['RawMessageDelivery'] == 'false':
+        return {}
+
+    message_attributes = {}
+    for key, value in attributes.items():
+        attribute = {}
+        attribute['DataType'] = value['Type']
+        if value['Type'] == 'Binary':
+            attribute['BinaryValue'] = value['Value']
+        else:
+            attribute['StringValue'] = value['Value']
+        message_attributes[key] = attribute
+
+    return message_attributes
+
+
 def get_message_attributes(req_data):
     attributes = {}
     x = 1

diff --git a/tests/unit/test_sns_listener.py b/tests/unit/test_sns_listener.py
@@ -120,6 +120,32 @@ def test_create_sns_message_body_json_structure_sqs_protocol(self):
 
         assert (result['Message'] == 'sqs message')
 
+    def test_create_sqs_message_attributes(self):
+        self.subscriber['RawMessageDelivery'] = 'true'
+        action = {
+            'Message': ['msg'],
+            'Subject': ['subject'],
+            'MessageAttributes.entry.1.Name': ['attr1'],
+            'MessageAttributes.entry.1.Value.DataType': ['String'],
+            'MessageAttributes.entry.1.Value.StringValue': ['value1'],
+            'MessageAttributes.entry.2.Name': ['attr2'],
+            'MessageAttributes.entry.2.Value.DataType': ['Binary'],
+            'MessageAttributes.entry.2.Value.BinaryValue': ['value2'.encode('utf-8')],
+            'MessageAttributes.entry.3.Name': ['attr3'],
+            'MessageAttributes.entry.3.Value.DataType': ['Number'],
+            'MessageAttributes.entry.3.Value.StringValue': ['value3'],
+        }
+
+        attributes = sns_listener.get_message_attributes(action)
+        result = sns_listener.create_sqs_message_attributes(self.subscriber, attributes)
+
+        assert (result['attr1']['DataType'] == 'String')
+        assert (result['attr1']['StringValue'] == 'value1')
+        assert (result['attr2']['DataType'] == 'Binary')
+        assert (result['attr2']['BinaryValue'] == 'value2'.encode('utf-8'))
+        assert (result['attr3']['DataType'] == 'Number')
+        assert (result['attr3']['StringValue'] == 'value3')
+
 
 def test_filter_policy():
     test_data = [