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)
async def message_producer(p: Producer) -> None:
""" Generate a new Position message, in various formats, and publish it """
while True:
msg = dict(latitude=130.0, longitude=-30.0, altitude=50.0)
content_type = random.choice([
CONTENT_TYPE_TEXT,
CONTENT_TYPE_JSON,
CONTENT_TYPE_MSGPACK,
CONTENT_TYPE_YAML,
CONTENT_TYPE_PROTOBUF,
])
compression_name = random.choice(
[None, COMPRESSION_GZIP, COMPRESSION_BZ2])
print(
f"Sending message using content-type={content_type}, compression={compression_name}"
)
if content_type == CONTENT_TYPE_TEXT:
msg = ",".join(f"{k}={v}" for k, v in msg.items())
elif content_type == CONTENT_TYPE_PROTOBUF:
msg = Position(**msg)
await p.publish_message(msg,
content_type=content_type,
compression=compression_name)
await asyncio.sleep(1.0)
async def message_producer(e) -> None:
""" Generate a new message and send it """
while True:
protobuf_data = Position(
latitude=130.0,
longitude=-30.0,
altitude=50.0,
status=Position.SIMULATED, # pylint: disable=no-member
)
print(f"sending a message: {protobuf_data}")
e.send(protobuf_data, type_identifier=type_identifier)
await asyncio.sleep(1.0)
def on_peer_available(cli: MtiStreamClient, peer_id):
print(f"Client {peer_id} connected")
# Upon connection, send a message to the server
protobuf_data = Position(
latitude=130.0,
longitude=-30.0,
altitude=50.0,
status=Position.SIMULATED, # pylint: disable=no-member
)
print(f"sending a message: {protobuf_data}")
client.send(protobuf_data,
peer_id=peer_id,
type_identifier=type_identifier)
async def test_topic_pubsub_protobuf(self):
""" check Protobuf messages can be published and received """
from position_pb2 import Position
exchange_name = "test"
p = Producer(
amqp_url=AMQP_URL,
exchange_name=exchange_name,
routing_key="position.update",
)
on_message_mock = unittest.mock.Mock()
c = Consumer(
amqp_url=AMQP_URL,
exchange_name=exchange_name,
routing_key="position.#",
on_message=on_message_mock,
)
# Register messages that require using the x-type-id message attribute
serializer = serialization.registry.get_serializer(
CONTENT_TYPE_PROTOBUF)
_type_identifier = serializer.registry.register_message(Position)
try:
await p.start()
await c.start()
await asyncio.sleep(0.1)
protobuf_msg = Position(**POSITION_DICT)
await p.publish_message(protobuf_msg,
content_type=CONTENT_TYPE_PROTOBUF)
await asyncio.sleep(0.1)
self.assertTrue(on_message_mock.called)
args, _kwargs = on_message_mock.call_args_list[0]
payload, message = args
self.assertEqual(payload, protobuf_msg)
self.assertEqual(message.properties.content_type,
CONTENT_TYPE_PROTOBUF)
self.assertEqual(message.headers.get("compression"), None)
finally:
await c.stop()
await p.stop()
async def on_message(server, data, peer_id, **kwargs) -> None:
print(f"Server received msg from {peer_id}: {data}")
# Wait briefly before sending a reply to the reply!
await asyncio.sleep(1)
protobuf_data = Position(
latitude=data.latitude + 0.1,
longitude=data.longitude + 0.1,
altitude=data.altitude + 0.1,
status=Position.SIMULATED, # pylint: disable=no-member
)
print(f"sending a message: {protobuf_data}")
server.send(protobuf_data,
peer_id=peer_id,
type_identifier=type_identifier)
def test_protobuf_serialization_roundtrip(self):
from position_pb2 import Position
protobuf_data = Position(latitude=130.0,
longitude=-30.0,
altitude=50.0,
status=Position.SIMULATED)
serializer = serialization.registry.get_serializer("protobuf")
type_identifier = serializer.registry.register_message(
Position, type_identifier=1)
content_type, content_encoding, payload = serialization.dumps(
protobuf_data, "protobuf", type_identifier=type_identifier)
self.assertIsInstance(payload, bytes)
self.assertEqual(content_type, serialization.CONTENT_TYPE_PROTOBUF)
self.assertEqual(content_encoding, "binary")
recovered_data = serialization.loads(
payload,
content_type=content_type,
content_encoding=content_encoding,
type_identifier=type_identifier,
)
self.assertEqual(protobuf_data, recovered_data)
async def test_protobuf_client_server_interaction_with_msg_id(self):
""" check protobuf client server interactions with message identifiers """
from position_pb2 import Position
protobuf_data = Position(latitude=130.0,
longitude=-30.0,
altitude=50.0,
status=Position.SIMULATED)
server_on_message_mock = asynctest.CoroutineMock()
server_on_started_mock = asynctest.CoroutineMock()
server_on_stopped_mock = asynctest.CoroutineMock()
server_on_peer_available_mock = asynctest.CoroutineMock()
server_on_peer_unavailable_mock = asynctest.CoroutineMock()
server_ep = MtiStreamServer(
on_message=server_on_message_mock,
on_started=server_on_started_mock,
on_stopped=server_on_stopped_mock,
on_peer_available=server_on_peer_available_mock,
on_peer_unavailable=server_on_peer_unavailable_mock,
content_type=serialization.CONTENT_TYPE_PROTOBUF,
)
await server_ep.start(addr="127.0.0.1", family=socket.AF_INET)
self.assertTrue(server_on_started_mock.called)
address, port = server_ep.bindings[0]
client_on_message_mock = asynctest.CoroutineMock()
client_on_started_mock = asynctest.CoroutineMock()
client_on_stopped_mock = asynctest.CoroutineMock()
client_on_peer_available_mock = asynctest.CoroutineMock()
client_on_peer_unavailable_mock = asynctest.CoroutineMock()
client_ep = MtiStreamClient(
on_message=client_on_message_mock,
on_started=client_on_started_mock,
on_stopped=client_on_stopped_mock,
on_peer_available=client_on_peer_available_mock,
on_peer_unavailable=client_on_peer_unavailable_mock,
content_type=serialization.CONTENT_TYPE_PROTOBUF,
)
await client_ep.start(addr=address, port=port, family=socket.AF_INET)
await asyncio.sleep(0.3)
self.assertTrue(client_on_started_mock.called)
self.assertTrue(client_on_peer_available_mock.called)
self.assertTrue(server_on_peer_available_mock.called)
self.assertEqual(len(client_ep.connections), 1)
# Register a message object with a unique message identifier. Only
# one of these calls is actually required to register the identifier
# with the object because a common serializer is used in this test
# case where both the client and server are instantiated.
# However, both calls are made to be more representative of how
# messages would be registered in a real application.
type_identifier = 1
server_ep.register_message(type_identifier, Position)
client_ep.register_message(type_identifier, Position)
test_msg_in = protobuf_data
# Send a msg with identifier from client to server
client_ep.send(test_msg_in, type_identifier=type_identifier)
await asyncio.sleep(0.1)
self.assertTrue(server_on_message_mock.called)
(args, kwargs) = server_on_message_mock.call_args_list[0]
_svr, received_msg = args
received_msg_id = kwargs["type_identifier"]
sender_id = kwargs["peer_id"]
self.assertEqual(received_msg, test_msg_in)
self.assertEqual(received_msg_id, type_identifier)
# Send a msg from server to client
server_ep.send(received_msg,
type_identifier=received_msg_id,
peer_id=sender_id)
await asyncio.sleep(0.1)
(args, kwargs) = client_on_message_mock.call_args_list[0]
_cli, received_msg = args
received_msg_id = kwargs["type_identifier"]
sender_id = kwargs["peer_id"]
self.assertEqual(received_msg, test_msg_in)
self.assertEqual(received_msg_id, type_identifier)
await client_ep.stop()
await asyncio.sleep(0.1)
self.assertTrue(client_on_stopped_mock.called)
self.assertTrue(client_on_peer_unavailable_mock.called)
self.assertTrue(server_on_peer_unavailable_mock.called)
await server_ep.stop()
self.assertTrue(server_on_stopped_mock.called)
async def test_protobuf_sender_receiver_interactions(self):
""" check protobuf sender and receiver interactions """
from position_pb2 import Position
protobuf_data = Position(latitude=130.0,
longitude=-30.0,
altitude=50.0,
status=Position.SIMULATED)
receiver_on_message_mock = asynctest.CoroutineMock()
receiver_on_started_mock = asynctest.CoroutineMock()
receiver_on_stopped_mock = asynctest.CoroutineMock()
receiver_on_peer_available_mock = asynctest.CoroutineMock()
receiver_on_peer_unavailable_mock = asynctest.CoroutineMock()
receiver_ep = MtiDatagramEndpoint(
on_message=receiver_on_message_mock,
on_started=receiver_on_started_mock,
on_stopped=receiver_on_stopped_mock,
on_peer_available=receiver_on_peer_available_mock,
on_peer_unavailable=receiver_on_peer_unavailable_mock,
content_type=serialization.CONTENT_TYPE_PROTOBUF,
)
await receiver_ep.start(local_addr=("127.0.0.1", 0))
self.assertTrue(receiver_on_started_mock.called)
address, port = receiver_ep.bindings[0]
sender_on_message_mock = asynctest.CoroutineMock()
sender_on_started_mock = asynctest.CoroutineMock()
sender_on_stopped_mock = asynctest.CoroutineMock()
sender_on_peer_available_mock = asynctest.CoroutineMock()
sender_on_peer_unavailable_mock = asynctest.CoroutineMock()
sender_ep = MtiDatagramEndpoint(
on_message=sender_on_message_mock,
on_started=sender_on_started_mock,
on_stopped=sender_on_stopped_mock,
on_peer_available=sender_on_peer_available_mock,
on_peer_unavailable=sender_on_peer_unavailable_mock,
content_type=serialization.CONTENT_TYPE_PROTOBUF,
)
await sender_ep.start(remote_addr=(address, port))
await asyncio.sleep(0.3)
self.assertTrue(sender_on_started_mock.called)
self.assertTrue(sender_on_peer_available_mock.called)
self.assertTrue(receiver_on_peer_available_mock.called)
type_identifier = 2
receiver_ep.register_message(type_identifier, Position)
# Send a msg with identifier from sender to receiver
test_msg_in = protobuf_data
sender_ep.send(test_msg_in, type_identifier=type_identifier)
await asyncio.sleep(0.1)
self.assertTrue(receiver_on_message_mock.called)
(args, kwargs) = receiver_on_message_mock.call_args_list[0]
ep, received_msg = args
_received_sender_id = kwargs["addr"]
received_msg_id = kwargs["type_identifier"]
self.assertIsInstance(ep, MtiDatagramEndpoint)
self.assertEqual(received_msg, test_msg_in)
self.assertEqual(received_msg_id, type_identifier)
await sender_ep.stop()
await asyncio.sleep(0.1)
self.assertTrue(sender_on_stopped_mock.called)
self.assertTrue(sender_on_peer_unavailable_mock.called)
await receiver_ep.stop()
self.assertTrue(receiver_on_stopped_mock.called)
self.assertTrue(receiver_on_peer_unavailable_mock.called)