From f9ce99f0d1ae79c27872d6e28a1ddb058a1eb113 Mon Sep 17 00:00:00 2001
From: usmangani1 <sgosman_chem@yahoo.com>
Date: Wed, 30 Sep 2020 23:10:00 +0530
Subject: [PATCH] Fix:SQS md5 calculation for custom string data type. (#3346)

* Fix:SQS md5 calculation for custom string data type.

* Linting

Co-authored-by: Bert Blommers <info@bertblommers.nl>
---
 moto/sqs/models.py         | 23 ++++++++++++++---------
 tests/test_sqs/test_sqs.py | 11 +++++++++++
 2 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/moto/sqs/models.py b/moto/sqs/models.py
index 720ab6e7..72218826 100644
--- a/moto/sqs/models.py
+++ b/moto/sqs/models.py
@@ -43,7 +43,12 @@ DEFAULT_SENDER_ID = "AIDAIT2UOQQY3AUEKVGXU"
 
 MAXIMUM_MESSAGE_LENGTH = 262144  # 256 KiB
 
-TRANSPORT_TYPE_ENCODINGS = {"String": b"\x01", "Binary": b"\x02", "Number": b"\x01"}
+TRANSPORT_TYPE_ENCODINGS = {
+    "String": b"\x01",
+    "Binary": b"\x02",
+    "Number": b"\x01",
+    "String.custom": b"\x01",
+}
 
 
 class Message(BaseModel):
@@ -88,14 +93,14 @@ class Message(BaseModel):
         struct_format = "!I".encode("ascii")  # ensure it's a bytestring
         for name in sorted(self.message_attributes.keys()):
             attr = self.message_attributes[name]
-            data_type_parts = attr["data_type"].split(".")
-            data_type = data_type_parts[0]
+            whole_data_type = attr.get("data_type")
+            if TRANSPORT_TYPE_ENCODINGS.get(whole_data_type):
+                data_type = whole_data_type
+            else:
+                data_type_parts = attr["data_type"].split(".")
+                data_type = data_type_parts[0]
 
-            if data_type not in [
-                "String",
-                "Binary",
-                "Number",
-            ]:
+            if data_type not in ["String", "Binary", "Number", "String.custom"]:
                 raise MessageAttributesInvalid(
                     "The message attribute '{0}' has an invalid message attribute type, the set of supported type prefixes is Binary, Number, and String.".format(
                         name[0]
@@ -112,7 +117,7 @@ class Message(BaseModel):
             encoded += struct.pack(struct_format, len(data_type)) + utf8(data_type)
             encoded += TRANSPORT_TYPE_ENCODINGS[data_type]
 
-            if data_type == "String" or data_type == "Number":
+            if data_type in ["String", "Number", "String.custom"]:
                 value = attr["string_value"]
             elif data_type == "Binary":
                 value = base64.b64decode(attr["binary_value"])
diff --git a/tests/test_sqs/test_sqs.py b/tests/test_sqs/test_sqs.py
index 9ce0f21c..48fa2029 100644
--- a/tests/test_sqs/test_sqs.py
+++ b/tests/test_sqs/test_sqs.py
@@ -658,6 +658,17 @@ def test_send_receive_message_with_attributes_with_labels():
         "994258b45346a2cc3f9cbb611aa7af30"
     )
 
+    response = queue.send_message(
+        MessageBody="test message",
+        MessageAttributes={
+            "somevalue": {"StringValue": "somevalue", "DataType": "String.custom",}
+        },
+    )
+
+    response.get("MD5OfMessageAttributes").should.equal(
+        "9e05cca738e70ff6c6041e82d5e77ef1"
+    )
+
 
 @mock_sqs
 def test_send_receive_message_timestamps():