py-libp2p/tests/pubsub/utils.py
2019-08-03 17:50:14 +08:00

83 lines
1.9 KiB
Python

import struct
from typing import Sequence
from libp2p.peer.id import ID
from libp2p.pubsub.pb import rpc_pb2
from tests.utils import connect
def message_id_generator(start_val):
"""
Generate a unique message id
:param start_val: value to start generating messages at
:return: message id
"""
val = start_val
def generator():
# Allow manipulation of val within closure
nonlocal val
# Increment id
val += 1
# Convert val to big endian
return struct.pack(">Q", val)
return generator
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)
)
# FIXME: There is no difference between `sparse_connect` and `dense_connect`,
# before `connect_some` is fixed.
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)
# 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
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)