py-libp2p/tests/pubsub/utils.py

83 lines
1.9 KiB
Python
Raw Normal View History

import struct
2019-08-01 06:00:12 +08:00
from typing import Sequence
from libp2p.peer.id import ID
from libp2p.pubsub.pb import rpc_pb2
2019-07-26 18:35:25 +08:00
from tests.utils import connect
def message_id_generator(start_val):
2019-04-02 04:55:44 +08:00
"""
Generate a unique message id
:param start_val: value to start generating messages at
:return: message id
2019-04-02 04:55:44 +08:00
"""
val = start_val
2019-07-26 18:35:25 +08:00
def generator():
# Allow manipulation of val within closure
nonlocal val
# Increment id
val += 1
# Convert val to big endian
2019-08-01 06:00:12 +08:00
return struct.pack(">Q", val)
return generator
2019-04-02 04:55:44 +08:00
def make_pubsub_msg(
2019-08-01 06:00:12 +08:00
origin_id: ID, topic_ids: Sequence[str], data: bytes, seqno: bytes
) -> rpc_pb2.Message:
return rpc_pb2.Message(
2019-08-01 06:00:12 +08:00
from_id=origin_id.to_bytes(), seqno=seqno, data=data, topicIDs=list(topic_ids)
)
2019-07-26 18:35:25 +08:00
# FIXME: There is no difference between `sparse_connect` and `dense_connect`,
# before `connect_some` is fixed.
2019-08-01 06:00:12 +08:00
async def sparse_connect(hosts):
await connect_some(hosts, 3)
async def dense_connect(hosts):
await connect_some(hosts, 10)
async def connect_all(hosts):
for i, host in enumerate(hosts):
for host2 in hosts[i + 1 :]:
await connect(host, host2)
2019-07-26 18:35:25 +08:00
# FIXME: `degree` is not used at all
async def connect_some(hosts, degree):
for i, host in enumerate(hosts):
for host2 in hosts[i + 1 :]:
await connect(host, host2)
# TODO: USE THE CODE BELOW
# for i, host in enumerate(hosts):
# j = 0
# while j < degree:
# n = random.randint(0, len(hosts) - 1)
# if n == i:
# j -= 1
# continue
# neighbor = hosts[n]
# await connect(host, neighbor)
# j += 1
2019-07-18 19:39:57 +08:00
2019-07-26 18:35:25 +08:00
2019-07-18 19:39:57 +08:00
async def one_to_all_connect(hosts, central_host_index):
for i, host in enumerate(hosts):
if i != central_host_index:
await connect(hosts[central_host_index], host)