def test_none_compression(self):
codecs = compression.registry.compressors
self.assertIn(None, codecs)
# check exception is raised when bytes are not passed in
with self.assertRaises(Exception) as cm:
compression.compress({})
self.assertIn("Can only compress bytes", str(cm.exception))
# perform roundtrip check
content_type, payload = compression.compress(TEST_DATA)
self.assertIsNone(content_type)
self.assertEqual(TEST_DATA, payload)
content_type, d = compression.decompress(payload)
self.assertIsNone(content_type)
self.assertEqual(d, TEST_DATA)
def test_snappy_compression(self):
convenience_name = "snappy"
mime_type = compression.COMPRESSION_SNAPPY
# check exception is raised when bytes are not passed in
with self.assertRaises(Exception) as cm:
compression.compress({}, mime_type)
self.assertIn("Can only compress bytes", str(cm.exception))
# perform roundtrip check
for c_name in (convenience_name, mime_type):
with self.subTest(f"Check snappy compression roundtrip using {c_name}"):
content_type, payload = compression.compress(TEST_DATA, c_name)
self.assertEqual(content_type, mime_type)
self.assertNotEqual(TEST_DATA, payload)
content_type, d = compression.decompress(payload, content_type)
self.assertEqual(content_type, mime_type)
self.assertEqual(d, TEST_DATA)
def test_compression_roundtrip(self):
codecs = compression.registry.compressors
for name, _settings in codecs.items():
with self.subTest(f"Check {name} compression roundtrip"):
convenience_name = name
mime_type = compression.registry.name_to_type[name]
# The convenience name or the content_type can be used when
# specifying compression method. Check both.
for c_label in (convenience_name, mime_type):
content_type, payload = compression.compress(TEST_DATA, c_label)
self.assertEqual(content_type, mime_type)
if c_label is None:
self.assertEqual(TEST_DATA, payload)
else:
self.assertNotEqual(TEST_DATA, payload)
content_type, d = compression.decompress(payload, content_type)
self.assertEqual(content_type, mime_type)
self.assertEqual(d, TEST_DATA)
def encode_payload(
data: Any,
*,
content_type: str = None,
compression: str = None,
headers: dict = None,
type_identifier: int = None,
) -> Tuple[bytes, Optional[str], str]:
""" Prepare a message payload.
:param data: The message data to encode.
:param content_type: A string specifying the message content type. By
default the value is None. This field determines the data
serialization format.
:param compression: An optional string specifying the compression strategy
to use. It can be provided using the convenience name or the mime-type.
If compression is defined then headers must also be supplied as
compression is passed as an attribute in message headers.
:param headers: A dict of headers that will be associated with the
message.
:param type_identifier: An integer that uniquely identifies a
registered message.
:returns: A three-item tuple containing the serialized data as bytes
a string specifying the content type (e.g., `application/json`) and
a string specifying the content encoding, (e.g. `utf-8`).
"""
# Some content-types require additional information to be passed to
# help decode the message payload. This is achieved by adding
# information to the message headers.
# Google Protocol buffer decoders require awareness of the object type
# being decoded (referred to as a symbol). The symbol id is added to
# the headers so that it can be used on the receiving side.
if content_type == CONTENT_TYPE_PROTOBUF:
serializer = registry.get_serializer(CONTENT_TYPE_PROTOBUF)
if not isinstance(headers, dict):
raise Exception("Headers must be supplied when using protobuf")
headers["x-type-id"] = serializer.registry.get_id_for_object(data)
# Avro decoders require awareness of the schema that describes the object.
# This information is added to the headers so that it can be used on the
# receiving side.
elif content_type == CONTENT_TYPE_AVRO:
if type_identifier is None:
raise Exception("No Avro id specified!")
if not isinstance(headers, dict):
raise Exception("Headers must be supplied when using Avro")
headers["x-type-id"] = type_identifier
serialization_name = registry.type_to_name[content_type]
try:
content_type, content_encoding, payload = dumps(
data, serialization_name, type_identifier=type_identifier
)
except Exception as exc:
raise Exception(f"Error serializing payload to {content_type}: {exc}") from None
if compression:
if not isinstance(headers, dict):
raise Exception("Headers must be supplied when using compression")
try:
headers["compression"], payload = compress(payload, compression)
except Exception as exc:
raise Exception(
f"Error compressing payload using {compression}: {exc}"
) from None
return payload, content_type, content_encoding
def test_compress_with_unspecified_name_or_type(self):
content_type, _payload = compression.compress(b"")
self.assertEqual(content_type, None)
def test_compress_with_invalid_name_or_type(self):
with self.assertRaises(Exception) as cm:
compression.compress(b"", "invalid")
self.assertIn("Invalid compressor", str(cm.exception))
