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>
This commit is contained in:
usmangani1 2020-09-30 23:10:00 +05:30 committed by GitHub
parent 1c3ba83fc2
commit f9ce99f0d1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 25 additions and 9 deletions

View File

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

View File

@ -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():