def test_encode_payload_exceptions(self):
TEXT_DATA = "The Quick Brown Fox Jumps Over The Lazy Dog"
with unittest.mock.patch("gestalt.amq.utils.dumps") as mock_dumps:
mock_dumps.side_effect = Exception("Boom!")
with self.assertRaises(Exception):
headers = {}
utils.encode_payload(
TEXT_DATA,
content_type=serialization.CONTENT_TYPE_TEXT,
compression=compression.COMPRESSION_ZLIB,
headers=headers,
)
with unittest.mock.patch(
"gestalt.amq.utils.compress") as mock_compress:
mock_compress.side_effect = Exception("Boom!")
with self.assertRaises(Exception):
headers = {}
utils.encode_payload(
TEXT_DATA,
content_type=serialization.CONTENT_TYPE_TEXT,
compression=compression.COMPRESSION_ZLIB,
headers=headers,
)
def test_json_payload_roundtrip(self):
JSON_DATA = dict(latitude=130.0, longitude=-30.0, altitude=50.0)
codecs = compression.registry.compressors
compression_names = list(codecs.keys())
for compression_name in compression_names:
with self.subTest(
f"Check compressing payload using {compression_name}"):
compression_mime_type = codecs[compression_name].content_type
headers = {}
payload, content_type, content_encoding = utils.encode_payload(
JSON_DATA,
content_type=serialization.CONTENT_TYPE_JSON,
compression=compression_name,
headers=headers,
)
if compression_name:
self.assertIn("compression", headers)
data = utils.decode_payload(
payload,
compression=compression_mime_type,
content_type=content_type,
content_encoding=content_encoding,
)
self.assertEqual(data, JSON_DATA)
def test_text_payload_roundtrip(self):
TEXT_DATA = "The Quick Brown Fox Jumps Over The Lazy Dog"
codecs = compression.registry.compressors
compression_names = list(codecs.keys())
for compression_name in compression_names:
with self.subTest(
f"Check compressing payload using {compression_name}"):
compression_mime_type = codecs[compression_name].content_type
headers = {}
payload, content_type, content_encoding = utils.encode_payload(
TEXT_DATA,
content_type=serialization.CONTENT_TYPE_TEXT,
compression=compression_name,
headers=headers,
)
if compression_name:
self.assertIn("compression", headers)
data = utils.decode_payload(
payload,
compression=compression_mime_type,
content_type=content_type,
content_encoding=content_encoding,
)
self.assertEqual(data, TEXT_DATA)
def test_protobuf_payload_roundtrip(self):
from position_pb2 import Position
PROTOBUF_DATA = Position(
latitude=130.0,
longitude=-30.0,
altitude=50.0,
status=Position.SIMULATED, # pylint: disable=no-member
)
serializer = serialization.registry.get_serializer(
serialization.CONTENT_TYPE_PROTOBUF)
_type_identifier = serializer.registry.register_message(Position)
# Protobuf is already a compact binary format so there is nothing
# to be gained by compressing it further.
headers = {}
payload, content_type, content_encoding = utils.encode_payload(
PROTOBUF_DATA,
content_type=serialization.CONTENT_TYPE_PROTOBUF,
headers=headers,
)
self.assertNotIn("compression", headers)
self.assertIn("x-type-id", headers)
data = utils.decode_payload(
payload,
content_type=content_type,
content_encoding=content_encoding,
type_identifier=headers["x-type-id"],
)
self.assertEqual(data, PROTOBUF_DATA)
def test_msgpack_payload_roundtrip(self):
MSGPACK_DATA = dict(latitude=130.0, longitude=-30.0, altitude=50.0)
# Msgpack is already a compact binary format so there is nothing to
# be gained by compressing it further.
headers = {}
payload, content_type, content_encoding = utils.encode_payload(
MSGPACK_DATA,
content_type=serialization.CONTENT_TYPE_MSGPACK,
headers=headers,
)
self.assertNotIn("compression", headers)
data = utils.decode_payload(payload,
content_type=content_type,
content_encoding=content_encoding)
self.assertEqual(data, MSGPACK_DATA)
async def publish_message(
self,
data: Any,
routing_key: str = None,
content_type: str = None,
compression: str = None,
headers: Dict = None,
type_identifier: int = None,
):
""" Publish a message.
:param data: The message data to transfer.
:param routing-key: A routing key to use when publishing the message.
The routing key determines which queues the exchange passes the
message to. If not specified then the default routing key provided
to the class initializer is used.
:param content_type: A string defining the message serialization
content type. If not specified then the default serialization
strategy will be assumed.
:param compression: The name of the compression strategy to use. The
default value is None. In this case the default compression strategy
will be used. If that is None then no compression is used.
:param headers: Arbitrary headers to pass along with message.
:param type_identifier: An integer that uniquely identifies a
registered message. This parameter is only needed for some
serialization methods that do not code in type awareness, such
as Avro and Protobuf.
"""
if self.connection is None:
logger.error("Producer does not have a connection")
return
routing_key = routing_key if routing_key else self.routing_key
serialization = content_type if content_type else self.serialization
compression = compression if compression else self.compression
headers = {}
try:
payload, content_type, content_encoding = utils.encode_payload(
data,
content_type=serialization,
compression=compression,
headers=headers,
type_identifier=type_identifier,
)
except Exception:
logger.exception("Error encoding payload")
return
assert self.exchange is not None
await self.exchange.publish(
Message(
body=payload,
content_type=content_type,
content_encoding=content_encoding,
timestamp=time.time(),
headers=headers,
),
routing_key=routing_key,
mandatory=False, # don't care if no routes are actively consuming
)
async def _on_request_message(self, message: IncomingMessage):
""" Process a request message """
if not message.reply_to:
logger.warning(
f"Received message without 'reply_to' header. {message}")
await message.ack()
return
try:
payload = utils.decode_message(message)
except Exception:
logger.exception("Problem in message decode function")
await message.reject(requeue=False)
return
try:
assert self._request_handler is not None
response = self._request_handler(payload, message)
if inspect.isawaitable(response):
response = await response # type: ignore
except Exception:
logger.exception(f"Problem in user message handler function")
await message.reject(requeue=False)
return
headers = {} # type: Dict[str, str]
try:
payload, content_type, content_encoding = utils.encode_payload(
response,
content_type=self.serialization,
compression=self.compression,
headers=headers,
)
except Exception:
logger.exception("Error encoding response payload")
await message.reject(requeue=False)
return
response_message = Message(
body=payload,
content_type=content_type,
content_encoding=content_encoding,
timestamp=time.time(),
correlation_id=message.correlation_id,
delivery_mode=message.delivery_mode, # type: ignore
headers=headers,
)
try:
# Responses should always use the default exchange to route the
# message to the queue created by the consumer.
assert self.channel is not None
await self.channel.default_exchange.publish(response_message,
message.reply_to,
mandatory=False)
except Exception:
logger.exception(f"Failed to send response: {response_message}")
await message.reject(requeue=False)
return
await message.ack()
async def request(
self,
data: Any,
*,
service_name: str = None,
expiration: int = None,
delivery_mode: DeliveryMode = DeliveryMode.NOT_PERSISTENT,
) -> asyncio.Future:
""" Send a message to a service and wait for a response.
:param data: The message data to transfer.
:param expiration: An optional value representing the number of
seconds that a message can remain in a queue before being returned
and a timeout exception (:class:`asyncio.TimeoutError`) is raised.
:param delivery_mode: Request message delivery mode. Default is
not-persistent.
:raises asyncio.TimeoutError: When message expires before being handled.
:raises CancelledError: when called :func:`RPC.cancel`
:raises Exception: internal error
"""
service_name = service_name if service_name else self.service_name
correlation_id, future = self.create_future()
headers = {} # type: Dict[str, str]
# An exception may be raised here if the message can not be serialized.
payload, content_type, content_encoding = utils.encode_payload(
data,
content_type=self.serialization,
compression=self.compression,
headers=headers,
)
assert self.response_queue is not None
# Add a 'From' entry to message headers which will be used to route an
# expired message to the dead letter exchange queue.
headers["From"] = self.response_queue.name
message = Message(
body=payload,
content_type=content_type,
content_encoding=content_encoding,
timestamp=time.time(),
correlation_id=correlation_id,
delivery_mode=delivery_mode,
reply_to=self.response_queue.name,
headers=headers,
)
if expiration is not None:
message.expiration = expiration
logger.debug(
f"Sending request to {service_name} with correlation_id: {correlation_id}"
)
assert self.exchange is not None
await self.exchange.publish(
message,
routing_key=service_name,
mandatory=True, # report error if no queues are actively consuming
)
logger.debug(f"Waiting for response from {service_name}")
return await future
def test_avro_payload_roundtrip(self):
# Add schema to serializer schema registry
message_schema = {
"namespace":
"unittest.test_amq_utils",
"type":
"record",
"name":
"Position",
"fields": [
{
"name": "latitude",
"type": ["float", "null"]
},
{
"name": "longitude",
"type": ["float", "null"]
},
{
"name": "altitude",
"type": ["float", "null"]
},
],
}
serializer = serialization.registry.get_serializer(
serialization.CONTENT_TYPE_AVRO)
type_identifier = serializer.registry.register_message(message_schema)
AVRO_DATA = dict(latitude=130.0, longitude=-30.0, altitude=50.0)
codecs = compression.registry.compressors
compression_names = list(codecs.keys())
# Avro is already a compact binary format so there may not be much
# reason to compress it further, but whatever...
for compression_name in compression_names:
with self.subTest(
f"Check compressing payload using {compression_name}"):
compression_mime_type = codecs[compression_name].content_type
# When content-type is AVRO a type identifier must be supplied.
# Confirm an exception is raised if it is not provided.
with self.assertRaises(Exception):
headers = {}
payload, content_type, content_encoding = utils.encode_payload(
AVRO_DATA,
content_type=serialization.CONTENT_TYPE_AVRO,
compression=compression_name,
headers=headers,
)
headers = {}
payload, content_type, content_encoding = utils.encode_payload(
AVRO_DATA,
content_type=serialization.CONTENT_TYPE_AVRO,
compression=compression_name,
headers=headers,
type_identifier=type_identifier,
)
if compression_name:
self.assertIn("compression", headers)
self.assertIn("x-type-id", headers)
data = utils.decode_payload(
payload,
compression=compression_mime_type,
content_type=content_type,
content_encoding=content_encoding,
type_identifier=headers["x-type-id"],
)
self.assertEqual(data, AVRO_DATA)