async def _read_pubsub_msg() -> None:
writer_closed_task = asyncio.ensure_future(writer.wait_closed())
while True:
done, pending = await asyncio.wait(
[read_varint_prefixed_bytes(reader), writer_closed_task],
return_when=asyncio.FIRST_COMPLETED,
)
done_tasks = tuple(done)
if writer.is_closing():
return
read_task = done_tasks[0]
# Sanity check
assert read_task._coro.__name__ == "read_varint_prefixed_bytes"
msg_bytes = read_task.result()
ps_msg = p2pd_pb2.PSMessage()
ps_msg.ParseFromString(msg_bytes)
# Fill in the message used in py-libp2p
msg = rpc_pb2.Message(
from_id=ps_msg.from_id,
data=ps_msg.data,
seqno=ps_msg.seqno,
topicIDs=ps_msg.topicIDs,
signature=ps_msg.signature,
key=ps_msg.key,
)
queue.put_nowait(msg)
def make_pubsub_msg(origin_id: ID, topic_ids: Sequence[str], data: bytes,
seqno: bytes) -> rpc_pb2.Message:
return rpc_pb2.Message(
from_id=origin_id.to_bytes(),
seqno=seqno,
data=data,
topicIDs=list(topic_ids),
)
async def test_bcc_receive_server_handle_beacon_attestations(receive_server):
attestation = Attestation.create()
encoded_attestation = ssz.encode(attestation)
msg = rpc_pb2.Message(
from_id=b"my_id",
seqno=b"\x00" * 8,
data=encoded_attestation,
topicIDs=[PUBSUB_TOPIC_BEACON_ATTESTATION],
)
assert attestation not in receive_server.unaggregated_attestation_pool
beacon_attestation_queue = receive_server.topic_msg_queues[
PUBSUB_TOPIC_BEACON_ATTESTATION]
# Wait for receive server to process the new attestation
await beacon_attestation_queue.put(msg)
await wait_all_messages_processed(beacon_attestation_queue)
# Check that attestation is put to attestation pool
assert attestation in receive_server.unaggregated_attestation_pool
# Put the attestation in the next block
block = SignedBeaconBlockFactory(
parent=receive_server.chain.get_canonical_head())
block = block.transform(("message", "body", "attestations"),
(attestation, ))
encoded_block = ssz.encode(block, SignedBeaconBlock)
msg = rpc_pb2.Message(
from_id=b"my_id",
seqno=b"\x00" * 8,
data=encoded_block,
topicIDs=[PUBSUB_TOPIC_BEACON_BLOCK],
)
beacon_block_queue = receive_server.topic_msg_queues[
PUBSUB_TOPIC_BEACON_BLOCK]
# Wait for receive server to process the new block
await beacon_block_queue.put(msg)
await wait_all_messages_processed(beacon_block_queue)
# Check that attestation is removed from attestation pool
assert attestation not in receive_server.unaggregated_attestation_pool
async def test_bcc_receive_server_handle_beacon_attestations(receive_server):
attestation = Attestation()
encoded_attestation = ssz.encode(attestation)
msg = rpc_pb2.Message(
from_id=b"my_id",
seqno=b"\x00" * 8,
data=encoded_attestation,
topicIDs=[PUBSUB_TOPIC_BEACON_ATTESTATION],
)
assert attestation not in receive_server.attestation_pool
beacon_attestation_queue = receive_server.topic_msg_queues[
PUBSUB_TOPIC_BEACON_ATTESTATION]
# Wait for receive server to process the new attestation
await beacon_attestation_queue.put(msg)
await wait_all_messages_processed(beacon_attestation_queue)
# Check that attestation is put to attestation pool
assert attestation in receive_server.attestation_pool
# Put the attestation in the next block
block = get_blocks(receive_server.chain, num_blocks=1)[0]
block = block.copy(body=block.body.copy(attestations=[attestation]))
encoded_block = ssz.encode(block, BeaconBlock)
msg = rpc_pb2.Message(
from_id=b"my_id",
seqno=b"\x00" * 8,
data=encoded_block,
topicIDs=[PUBSUB_TOPIC_BEACON_BLOCK],
)
beacon_block_queue = receive_server.topic_msg_queues[
PUBSUB_TOPIC_BEACON_BLOCK]
# Wait for receive server to process the new block
await beacon_block_queue.put(msg)
await wait_all_messages_processed(beacon_block_queue)
# Check that attestation is removed from attestation pool
assert attestation not in receive_server.attestation_pool
async def _read_pubsub_msg() -> None:
while True:
msg_bytes = await read_varint_prefixed_bytes(stream)
ps_msg = p2pd_pb2.PSMessage()
ps_msg.ParseFromString(msg_bytes)
# Fill in the message used in py-libp2p
msg = rpc_pb2.Message(
from_id=ps_msg.from_id,
data=ps_msg.data,
seqno=ps_msg.seqno,
topicIDs=ps_msg.topicIDs,
signature=ps_msg.signature,
key=ps_msg.key,
)
await send_channel.send(msg)
async def test_bcc_receive_server_handle_beacon_blocks(receive_server):
block = SignedBeaconBlockFactory(parent=receive_server.chain.get_canonical_head())
encoded_block = ssz.encode(block, SignedBeaconBlock)
msg = rpc_pb2.Message(
from_id=b"my_id",
seqno=b"\x00" * 8,
data=encoded_block,
topicIDs=[PUBSUB_TOPIC_BEACON_BLOCK],
)
assert receive_server.chain.get_canonical_head() != block
beacon_block_queue = receive_server.topic_msg_queues[PUBSUB_TOPIC_BEACON_BLOCK]
await beacon_block_queue.put(msg)
await wait_all_messages_processed(beacon_block_queue)
assert receive_server.chain.get_canonical_head() == block
async def test_bcc_receive_server_handle_beacon_blocks(receive_server):
block = get_blocks(receive_server.chain, num_blocks=1)[0]
encoded_block = ssz.encode(block, BeaconBlock)
msg = rpc_pb2.Message(
from_id=b"my_id",
seqno=b"\x00" * 8,
data=encoded_block,
topicIDs=[PUBSUB_TOPIC_BEACON_BLOCK],
)
assert receive_server.chain.get_canonical_head() != block
beacon_block_queue = receive_server.topic_msg_queues[
PUBSUB_TOPIC_BEACON_BLOCK]
await beacon_block_queue.put(msg)
# Wait for receive server to process the new block
await asyncio.sleep(0.5)
assert receive_server.chain.get_canonical_head() == block
def generate_RPC_packet(origin_id, topics, msg_content, msg_id):
"""
Generate RPC packet to send over wire
:param origin_id: peer id of the message origin
:param topics: list of topics
:param msg_content: string of content in data
:param msg_id: seqno for the message
"""
packet = rpc_pb2.RPC()
message = rpc_pb2.Message(
from_id=origin_id.encode('utf-8'),
seqno=msg_id,
data=msg_content.encode('utf-8'),
)
for topic in topics:
message.topicIDs.extend([topic.encode('utf-8')])
packet.publish.extend([message])
return packet
async def test_continuously_read_stream(pubsubs_fsub, monkeypatch):
stream = FakeNetStream()
await pubsubs_fsub[0].subscribe(TESTING_TOPIC)
event_push_msg = asyncio.Event()
event_handle_subscription = asyncio.Event()
event_handle_rpc = asyncio.Event()
async def mock_push_msg(msg_forwarder, msg):
event_push_msg.set()
def mock_handle_subscription(origin_id, sub_message):
event_handle_subscription.set()
async def mock_handle_rpc(rpc, sender_peer_id):
event_handle_rpc.set()
monkeypatch.setattr(pubsubs_fsub[0], "push_msg", mock_push_msg)
monkeypatch.setattr(
pubsubs_fsub[0], "handle_subscription", mock_handle_subscription
)
monkeypatch.setattr(pubsubs_fsub[0].router, "handle_rpc", mock_handle_rpc)
async def wait_for_event_occurring(event):
try:
await asyncio.wait_for(event.wait(), timeout=1)
except asyncio.TimeoutError as error:
event.clear()
raise asyncio.TimeoutError(
f"Event {event} is not set before the timeout. "
"This indicates the mocked functions are not called properly."
) from error
else:
event.clear()
# Kick off the task `continuously_read_stream`
task = asyncio.ensure_future(pubsubs_fsub[0].continuously_read_stream(stream))
# Test: `push_msg` is called when publishing to a subscribed topic.
publish_subscribed_topic = rpc_pb2.RPC(
publish=[rpc_pb2.Message(topicIDs=[TESTING_TOPIC])]
)
await stream.write(
encode_varint_prefixed(publish_subscribed_topic.SerializeToString())
)
await wait_for_event_occurring(event_push_msg)
# Make sure the other events are not emitted.
with pytest.raises(asyncio.TimeoutError):
await wait_for_event_occurring(event_handle_subscription)
with pytest.raises(asyncio.TimeoutError):
await wait_for_event_occurring(event_handle_rpc)
# Test: `push_msg` is not called when publishing to a topic-not-subscribed.
publish_not_subscribed_topic = rpc_pb2.RPC(
publish=[rpc_pb2.Message(topicIDs=["NOT_SUBSCRIBED"])]
)
await stream.write(
encode_varint_prefixed(publish_not_subscribed_topic.SerializeToString())
)
with pytest.raises(asyncio.TimeoutError):
await wait_for_event_occurring(event_push_msg)
# Test: `handle_subscription` is called when a subscription message is received.
subscription_msg = rpc_pb2.RPC(subscriptions=[rpc_pb2.RPC.SubOpts()])
await stream.write(encode_varint_prefixed(subscription_msg.SerializeToString()))
await wait_for_event_occurring(event_handle_subscription)
# Make sure the other events are not emitted.
with pytest.raises(asyncio.TimeoutError):
await wait_for_event_occurring(event_push_msg)
with pytest.raises(asyncio.TimeoutError):
await wait_for_event_occurring(event_handle_rpc)
# Test: `handle_rpc` is called when a control message is received.
control_msg = rpc_pb2.RPC(control=rpc_pb2.ControlMessage())
await stream.write(encode_varint_prefixed(control_msg.SerializeToString()))
await wait_for_event_occurring(event_handle_rpc)
# Make sure the other events are not emitted.
with pytest.raises(asyncio.TimeoutError):
await wait_for_event_occurring(event_push_msg)
with pytest.raises(asyncio.TimeoutError):
await wait_for_event_occurring(event_handle_subscription)
task.cancel()
async def test_continuously_read_stream(monkeypatch, nursery,
security_protocol):
async def wait_for_event_occurring(event):
await trio.hazmat.checkpoint()
with trio.fail_after(0.1):
await event.wait()
class Events(NamedTuple):
push_msg: trio.Event
handle_subscription: trio.Event
handle_rpc: trio.Event
@contextmanager
def mock_methods():
event_push_msg = trio.Event()
event_handle_subscription = trio.Event()
event_handle_rpc = trio.Event()
async def mock_push_msg(msg_forwarder, msg):
event_push_msg.set()
await trio.hazmat.checkpoint()
def mock_handle_subscription(origin_id, sub_message):
event_handle_subscription.set()
async def mock_handle_rpc(rpc, sender_peer_id):
event_handle_rpc.set()
await trio.hazmat.checkpoint()
with monkeypatch.context() as m:
m.setattr(pubsubs_fsub[0], "push_msg", mock_push_msg)
m.setattr(pubsubs_fsub[0], "handle_subscription",
mock_handle_subscription)
m.setattr(pubsubs_fsub[0].router, "handle_rpc", mock_handle_rpc)
yield Events(event_push_msg, event_handle_subscription,
event_handle_rpc)
async with PubsubFactory.create_batch_with_floodsub(
1, security_protocol=security_protocol
) as pubsubs_fsub, net_stream_pair_factory(
security_protocol=security_protocol) as stream_pair:
await pubsubs_fsub[0].subscribe(TESTING_TOPIC)
# Kick off the task `continuously_read_stream`
nursery.start_soon(pubsubs_fsub[0].continuously_read_stream,
stream_pair[0])
# Test: `push_msg` is called when publishing to a subscribed topic.
publish_subscribed_topic = rpc_pb2.RPC(
publish=[rpc_pb2.Message(topicIDs=[TESTING_TOPIC])])
with mock_methods() as events:
await stream_pair[1].write(
encode_varint_prefixed(
publish_subscribed_topic.SerializeToString()))
await wait_for_event_occurring(events.push_msg)
# Make sure the other events are not emitted.
with pytest.raises(trio.TooSlowError):
await wait_for_event_occurring(events.handle_subscription)
with pytest.raises(trio.TooSlowError):
await wait_for_event_occurring(events.handle_rpc)
# Test: `push_msg` is not called when publishing to a topic-not-subscribed.
publish_not_subscribed_topic = rpc_pb2.RPC(
publish=[rpc_pb2.Message(topicIDs=["NOT_SUBSCRIBED"])])
with mock_methods() as events:
await stream_pair[1].write(
encode_varint_prefixed(
publish_not_subscribed_topic.SerializeToString()))
with pytest.raises(trio.TooSlowError):
await wait_for_event_occurring(events.push_msg)
# Test: `handle_subscription` is called when a subscription message is received.
subscription_msg = rpc_pb2.RPC(subscriptions=[rpc_pb2.RPC.SubOpts()])
with mock_methods() as events:
await stream_pair[1].write(
encode_varint_prefixed(subscription_msg.SerializeToString()))
await wait_for_event_occurring(events.handle_subscription)
# Make sure the other events are not emitted.
with pytest.raises(trio.TooSlowError):
await wait_for_event_occurring(events.push_msg)
with pytest.raises(trio.TooSlowError):
await wait_for_event_occurring(events.handle_rpc)
# Test: `handle_rpc` is called when a control message is received.
control_msg = rpc_pb2.RPC(control=rpc_pb2.ControlMessage())
with mock_methods() as events:
await stream_pair[1].write(
encode_varint_prefixed(control_msg.SerializeToString()))
await wait_for_event_occurring(events.handle_rpc)
# Make sure the other events are not emitted.
with pytest.raises(trio.TooSlowError):
await wait_for_event_occurring(events.push_msg)
with pytest.raises(trio.TooSlowError):
await wait_for_event_occurring(events.handle_subscription)
def make_pubsub_msg():
return rpc_pb2.Message()
