mirror of
https://github.com/varun-r-mallya/py-libp2p.git
synced 2026-02-12 16:10:57 +00:00
Fix all modules except for security
This commit is contained in:
mhchia
@ -1,30 +1,29 @@
|
||||
from contextlib import AsyncExitStack, asynccontextmanager
|
||||
from typing import Any, AsyncIterator, Dict, Tuple, cast
|
||||
from typing import Any, AsyncIterator, Dict, Sequence, Tuple, cast
|
||||
|
||||
from async_service import background_trio_service
|
||||
import factory
|
||||
import trio
|
||||
|
||||
from libp2p.tools.constants import GOSSIPSUB_PARAMS
|
||||
from libp2p import generate_new_rsa_identity, generate_peer_id_from
|
||||
from libp2p.crypto.keys import KeyPair
|
||||
from libp2p.host.basic_host import BasicHost
|
||||
from libp2p.host.routed_host import RoutedHost
|
||||
from libp2p.tools.utils import set_up_routers
|
||||
from libp2p.kademlia.network import KademliaServer
|
||||
from libp2p.host.host_interface import IHost
|
||||
from libp2p.network.connection.swarm_connection import SwarmConn
|
||||
from libp2p.network.stream.net_stream_interface import INetStream
|
||||
from libp2p.network.swarm import Swarm
|
||||
from libp2p.peer.peerstore import PeerStore
|
||||
from libp2p.peer.id import ID
|
||||
from libp2p.peer.peerstore import PeerStore
|
||||
from libp2p.pubsub.floodsub import FloodSub
|
||||
from libp2p.pubsub.gossipsub import GossipSub
|
||||
from libp2p.pubsub.pubsub import Pubsub
|
||||
from libp2p.pubsub.pubsub_router_interface import IPubsubRouter
|
||||
from libp2p.security.base_transport import BaseSecureTransport
|
||||
from libp2p.security.insecure.transport import PLAINTEXT_PROTOCOL_ID, InsecureTransport
|
||||
import libp2p.security.secio.transport as secio
|
||||
from libp2p.stream_muxer.mplex.mplex import MPLEX_PROTOCOL_ID, Mplex
|
||||
from libp2p.stream_muxer.mplex.mplex_stream import MplexStream
|
||||
from libp2p.tools.constants import GOSSIPSUB_PARAMS
|
||||
from libp2p.transport.tcp.tcp import TCP
|
||||
from libp2p.transport.typing import TMuxerOptions
|
||||
from libp2p.transport.upgrader import TransportUpgrader
|
||||
@ -74,7 +73,7 @@ class SwarmFactory(factory.Factory):
|
||||
@asynccontextmanager
|
||||
async def create_and_listen(
|
||||
cls, is_secure: bool, key_pair: KeyPair = None, muxer_opt: TMuxerOptions = None
|
||||
) -> Swarm:
|
||||
) -> AsyncIterator[Swarm]:
|
||||
# `factory.Factory.__init__` does *not* prepare a *default value* if we pass
|
||||
# an argument explicitly with `None`. If an argument is `None`, we don't pass it to
|
||||
# `factory.Factory.__init__`, in order to let the function initialize it.
|
||||
@ -92,7 +91,7 @@ class SwarmFactory(factory.Factory):
|
||||
@asynccontextmanager
|
||||
async def create_batch_and_listen(
|
||||
cls, is_secure: bool, number: int, muxer_opt: TMuxerOptions = None
|
||||
) -> Tuple[Swarm, ...]:
|
||||
) -> AsyncIterator[Tuple[Swarm, ...]]:
|
||||
async with AsyncExitStack() as stack:
|
||||
ctx_mgrs = [
|
||||
await stack.enter_async_context(
|
||||
@ -100,7 +99,7 @@ class SwarmFactory(factory.Factory):
|
||||
)
|
||||
for _ in range(number)
|
||||
]
|
||||
yield ctx_mgrs
|
||||
yield tuple(ctx_mgrs)
|
||||
|
||||
|
||||
class HostFactory(factory.Factory):
|
||||
@ -120,7 +119,7 @@ class HostFactory(factory.Factory):
|
||||
@asynccontextmanager
|
||||
async def create_batch_and_listen(
|
||||
cls, is_secure: bool, number: int
|
||||
) -> Tuple[BasicHost, ...]:
|
||||
) -> AsyncIterator[Tuple[BasicHost, ...]]:
|
||||
key_pairs = [generate_new_rsa_identity() for _ in range(number)]
|
||||
async with AsyncExitStack() as stack:
|
||||
swarms = [
|
||||
@ -136,30 +135,6 @@ class HostFactory(factory.Factory):
|
||||
yield hosts
|
||||
|
||||
|
||||
class RoutedHostFactory(factory.Factory):
|
||||
class Meta:
|
||||
model = RoutedHost
|
||||
|
||||
public_key = factory.LazyAttribute(lambda o: o.key_pair.public_key)
|
||||
network = factory.LazyAttribute(
|
||||
lambda o: SwarmFactory(is_secure=o.is_secure, key_pair=o.key_pair)
|
||||
)
|
||||
router = factory.LazyFunction(KademliaServer)
|
||||
|
||||
@classmethod
|
||||
@asynccontextmanager
|
||||
async def create_batch_and_listen(
|
||||
cls, is_secure: bool, number: int
|
||||
) -> Tuple[RoutedHost, ...]:
|
||||
key_pairs = [generate_new_rsa_identity() for _ in range(number)]
|
||||
routers = await set_up_routers((0,) * number)
|
||||
async with SwarmFactory.create_batch_and_listen(is_secure, number) as swarms:
|
||||
yield tuple(
|
||||
RoutedHost(key_pair.public_key, swarm, router)
|
||||
for key_pair, swarm, router in zip(key_pairs, swarms, routers)
|
||||
)
|
||||
|
||||
|
||||
class FloodsubFactory(factory.Factory):
|
||||
class Meta:
|
||||
model = FloodSub
|
||||
@ -191,17 +166,22 @@ class PubsubFactory(factory.Factory):
|
||||
|
||||
@classmethod
|
||||
@asynccontextmanager
|
||||
async def create_and_start(cls, host, router, cache_size):
|
||||
async def create_and_start(
|
||||
cls, host: IHost, router: IPubsubRouter, cache_size: int
|
||||
) -> AsyncIterator[Pubsub]:
|
||||
pubsub = PubsubFactory(host=host, router=router, cache_size=cache_size)
|
||||
async with background_trio_service(pubsub):
|
||||
yield pubsub
|
||||
|
||||
@classmethod
|
||||
@asynccontextmanager
|
||||
async def create_batch_with_floodsub(
|
||||
cls, number: int, is_secure: bool = False, cache_size: int = None
|
||||
):
|
||||
floodsubs = FloodsubFactory.create_batch(number)
|
||||
async def _create_batch_with_router(
|
||||
cls,
|
||||
number: int,
|
||||
routers: Sequence[IPubsubRouter],
|
||||
is_secure: bool = False,
|
||||
cache_size: int = None,
|
||||
) -> AsyncIterator[Tuple[Pubsub, ...]]:
|
||||
async with HostFactory.create_batch_and_listen(is_secure, number) as hosts:
|
||||
# Pubsubs should exit before hosts
|
||||
async with AsyncExitStack() as stack:
|
||||
@ -209,21 +189,80 @@ class PubsubFactory(factory.Factory):
|
||||
await stack.enter_async_context(
|
||||
cls.create_and_start(host, router, cache_size)
|
||||
)
|
||||
for host, router in zip(hosts, floodsubs)
|
||||
for host, router in zip(hosts, routers)
|
||||
]
|
||||
yield pubsubs
|
||||
yield tuple(pubsubs)
|
||||
|
||||
# @classmethod
|
||||
# async def create_batch_with_gossipsub(
|
||||
# cls, number: int, cache_size: int = None, gossipsub_params=GOSSIPSUB_PARAMS
|
||||
# ):
|
||||
# ...
|
||||
@classmethod
|
||||
@asynccontextmanager
|
||||
async def create_batch_with_floodsub(
|
||||
cls,
|
||||
number: int,
|
||||
is_secure: bool = False,
|
||||
cache_size: int = None,
|
||||
protocols: Sequence[TProtocol] = None,
|
||||
) -> AsyncIterator[Tuple[Pubsub, ...]]:
|
||||
if protocols is not None:
|
||||
floodsubs = FloodsubFactory.create_batch(number, protocols=list(protocols))
|
||||
else:
|
||||
floodsubs = FloodsubFactory.create_batch(number)
|
||||
async with cls._create_batch_with_router(
|
||||
number, floodsubs, is_secure, cache_size
|
||||
) as pubsubs:
|
||||
yield pubsubs
|
||||
|
||||
@classmethod
|
||||
@asynccontextmanager
|
||||
async def create_batch_with_gossipsub(
|
||||
cls,
|
||||
number: int,
|
||||
*,
|
||||
is_secure: bool = False,
|
||||
cache_size: int = None,
|
||||
protocols: Sequence[TProtocol] = None,
|
||||
degree: int = GOSSIPSUB_PARAMS.degree,
|
||||
degree_low: int = GOSSIPSUB_PARAMS.degree_low,
|
||||
degree_high: int = GOSSIPSUB_PARAMS.degree_high,
|
||||
time_to_live: int = GOSSIPSUB_PARAMS.time_to_live,
|
||||
gossip_window: int = GOSSIPSUB_PARAMS.gossip_window,
|
||||
gossip_history: int = GOSSIPSUB_PARAMS.gossip_history,
|
||||
heartbeat_interval: float = GOSSIPSUB_PARAMS.heartbeat_interval,
|
||||
) -> AsyncIterator[Tuple[Pubsub, ...]]:
|
||||
if protocols is not None:
|
||||
gossipsubs = GossipsubFactory.create_batch(
|
||||
number,
|
||||
protocols=protocols,
|
||||
degree=degree,
|
||||
degree_low=degree_low,
|
||||
degree_high=degree_high,
|
||||
time_to_live=time_to_live,
|
||||
gossip_window=gossip_window,
|
||||
heartbeat_interval=heartbeat_interval,
|
||||
)
|
||||
else:
|
||||
gossipsubs = GossipsubFactory.create_batch(
|
||||
number,
|
||||
degree=degree,
|
||||
degree_low=degree_low,
|
||||
degree_high=degree_high,
|
||||
time_to_live=time_to_live,
|
||||
gossip_window=gossip_window,
|
||||
heartbeat_interval=heartbeat_interval,
|
||||
)
|
||||
|
||||
async with cls._create_batch_with_router(
|
||||
number, gossipsubs, is_secure, cache_size
|
||||
) as pubsubs:
|
||||
async with AsyncExitStack() as stack:
|
||||
for router in gossipsubs:
|
||||
await stack.enter_async_context(background_trio_service(router))
|
||||
yield pubsubs
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def swarm_pair_factory(
|
||||
is_secure: bool, muxer_opt: TMuxerOptions = None
|
||||
) -> Tuple[Swarm, Swarm]:
|
||||
) -> AsyncIterator[Tuple[Swarm, Swarm]]:
|
||||
async with SwarmFactory.create_batch_and_listen(
|
||||
is_secure, 2, muxer_opt=muxer_opt
|
||||
) as swarms:
|
||||
@ -232,7 +271,9 @@ async def swarm_pair_factory(
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def host_pair_factory(is_secure: bool) -> Tuple[BasicHost, BasicHost]:
|
||||
async def host_pair_factory(
|
||||
is_secure: bool
|
||||
) -> AsyncIterator[Tuple[BasicHost, BasicHost]]:
|
||||
async with HostFactory.create_batch_and_listen(is_secure, 2) as hosts:
|
||||
await connect(hosts[0], hosts[1])
|
||||
yield hosts[0], hosts[1]
|
||||
@ -241,7 +282,7 @@ async def host_pair_factory(is_secure: bool) -> Tuple[BasicHost, BasicHost]:
|
||||
@asynccontextmanager
|
||||
async def swarm_conn_pair_factory(
|
||||
is_secure: bool, muxer_opt: TMuxerOptions = None
|
||||
) -> Tuple[SwarmConn, SwarmConn]:
|
||||
) -> AsyncIterator[Tuple[SwarmConn, SwarmConn]]:
|
||||
async with swarm_pair_factory(is_secure) as swarms:
|
||||
conn_0 = swarms[0].connections[swarms[1].get_peer_id()]
|
||||
conn_1 = swarms[1].connections[swarms[0].get_peer_id()]
|
||||
@ -249,7 +290,9 @@ async def swarm_conn_pair_factory(
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def mplex_conn_pair_factory(is_secure: bool) -> Tuple[Mplex, Mplex]:
|
||||
async def mplex_conn_pair_factory(
|
||||
is_secure: bool
|
||||
) -> AsyncIterator[Tuple[Mplex, Mplex]]:
|
||||
muxer_opt = {MPLEX_PROTOCOL_ID: Mplex}
|
||||
async with swarm_conn_pair_factory(is_secure, muxer_opt=muxer_opt) as swarm_pair:
|
||||
yield (
|
||||
@ -259,21 +302,25 @@ async def mplex_conn_pair_factory(is_secure: bool) -> Tuple[Mplex, Mplex]:
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def mplex_stream_pair_factory(is_secure: bool) -> Tuple[MplexStream, MplexStream]:
|
||||
async def mplex_stream_pair_factory(
|
||||
is_secure: bool
|
||||
) -> AsyncIterator[Tuple[MplexStream, MplexStream]]:
|
||||
async with mplex_conn_pair_factory(is_secure) as mplex_conn_pair_info:
|
||||
mplex_conn_0, mplex_conn_1 = mplex_conn_pair_info
|
||||
stream_0 = await mplex_conn_0.open_stream()
|
||||
stream_0 = cast(MplexStream, await mplex_conn_0.open_stream())
|
||||
await trio.sleep(0.01)
|
||||
stream_1: MplexStream
|
||||
async with mplex_conn_1.streams_lock:
|
||||
if len(mplex_conn_1.streams) != 1:
|
||||
raise Exception("Mplex should not have any other stream")
|
||||
stream_1 = tuple(mplex_conn_1.streams.values())[0]
|
||||
yield cast(MplexStream, stream_0), cast(MplexStream, stream_1)
|
||||
yield stream_0, stream_1
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def net_stream_pair_factory(is_secure: bool) -> Tuple[INetStream, INetStream]:
|
||||
async def net_stream_pair_factory(
|
||||
is_secure: bool
|
||||
) -> AsyncIterator[Tuple[INetStream, INetStream]]:
|
||||
protocol_id = TProtocol("/example/id/1")
|
||||
|
||||
stream_1: INetStream
|
||||
|
||||
@ -1,12 +1,11 @@
|
||||
import asyncio
|
||||
from typing import Dict
|
||||
import uuid
|
||||
from contextlib import AsyncExitStack, asynccontextmanager
|
||||
from typing import AsyncIterator, Dict, Tuple
|
||||
|
||||
from async_service import Service, background_trio_service
|
||||
|
||||
from libp2p.host.host_interface import IHost
|
||||
from libp2p.pubsub.floodsub import FloodSub
|
||||
from libp2p.pubsub.pubsub import Pubsub
|
||||
from libp2p.tools.constants import LISTEN_MADDR
|
||||
from libp2p.tools.factories import FloodsubFactory, PubsubFactory
|
||||
from libp2p.tools.factories import PubsubFactory
|
||||
|
||||
CRYPTO_TOPIC = "ethereum"
|
||||
|
||||
@ -18,7 +17,7 @@ CRYPTO_TOPIC = "ethereum"
|
||||
# Determine message type by looking at first item before first comma
|
||||
|
||||
|
||||
class DummyAccountNode:
|
||||
class DummyAccountNode(Service):
|
||||
"""
|
||||
Node which has an internal balance mapping, meant to serve as a dummy
|
||||
crypto blockchain.
|
||||
@ -27,19 +26,24 @@ class DummyAccountNode:
|
||||
crypto each user in the mappings holds
|
||||
"""
|
||||
|
||||
libp2p_node: IHost
|
||||
pubsub: Pubsub
|
||||
floodsub: FloodSub
|
||||
|
||||
def __init__(self, libp2p_node: IHost, pubsub: Pubsub, floodsub: FloodSub):
|
||||
self.libp2p_node = libp2p_node
|
||||
def __init__(self, pubsub: Pubsub) -> None:
|
||||
self.pubsub = pubsub
|
||||
self.floodsub = floodsub
|
||||
self.balances: Dict[str, int] = {}
|
||||
self.node_id = str(uuid.uuid1())
|
||||
|
||||
@property
|
||||
def host(self) -> IHost:
|
||||
return self.pubsub.host
|
||||
|
||||
async def run(self) -> None:
|
||||
self.subscription = await self.pubsub.subscribe(CRYPTO_TOPIC)
|
||||
self.manager.run_daemon_task(self.handle_incoming_msgs)
|
||||
await self.manager.wait_finished()
|
||||
|
||||
@classmethod
|
||||
async def create(cls) -> "DummyAccountNode":
|
||||
@asynccontextmanager
|
||||
async def create(cls, number: int) -> AsyncIterator[Tuple["DummyAccountNode", ...]]:
|
||||
"""
|
||||
Create a new DummyAccountNode and attach a libp2p node, a floodsub, and
|
||||
a pubsub instance to this new node.
|
||||
@ -47,15 +51,17 @@ class DummyAccountNode:
|
||||
We use create as this serves as a factory function and allows us
|
||||
to use async await, unlike the init function
|
||||
"""
|
||||
|
||||
pubsub = PubsubFactory(router=FloodsubFactory())
|
||||
await pubsub.host.get_network().listen(LISTEN_MADDR)
|
||||
return cls(libp2p_node=pubsub.host, pubsub=pubsub, floodsub=pubsub.router)
|
||||
async with PubsubFactory.create_batch_with_floodsub(number) as pubsubs:
|
||||
async with AsyncExitStack() as stack:
|
||||
dummy_acount_nodes = tuple(cls(pubsub) for pubsub in pubsubs)
|
||||
for node in dummy_acount_nodes:
|
||||
await stack.enter_async_context(background_trio_service(node))
|
||||
yield dummy_acount_nodes
|
||||
|
||||
async def handle_incoming_msgs(self) -> None:
|
||||
"""Handle all incoming messages on the CRYPTO_TOPIC from peers."""
|
||||
while True:
|
||||
incoming = await self.q.get()
|
||||
incoming = await self.subscription.receive()
|
||||
msg_comps = incoming.data.decode("utf-8").split(",")
|
||||
|
||||
if msg_comps[0] == "send":
|
||||
@ -63,13 +69,6 @@ class DummyAccountNode:
|
||||
elif msg_comps[0] == "set":
|
||||
self.handle_set_crypto(msg_comps[1], int(msg_comps[2]))
|
||||
|
||||
async def setup_crypto_networking(self) -> None:
|
||||
"""Subscribe to CRYPTO_TOPIC and perform call to function that handles
|
||||
all incoming messages on said topic."""
|
||||
self.q = await self.pubsub.subscribe(CRYPTO_TOPIC)
|
||||
|
||||
asyncio.ensure_future(self.handle_incoming_msgs())
|
||||
|
||||
async def publish_send_crypto(
|
||||
self, source_user: str, dest_user: str, amount: int
|
||||
) -> None:
|
||||
|
||||
@ -1,12 +1,10 @@
|
||||
# type: ignore
|
||||
# To add typing to this module, it's better to do it after refactoring test cases into classes
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
import trio
|
||||
|
||||
from libp2p.tools.constants import FLOODSUB_PROTOCOL_ID, LISTEN_MADDR
|
||||
from libp2p.tools.factories import PubsubFactory
|
||||
from libp2p.tools.constants import FLOODSUB_PROTOCOL_ID
|
||||
from libp2p.tools.utils import connect
|
||||
|
||||
SUPPORTED_PROTOCOLS = [FLOODSUB_PROTOCOL_ID]
|
||||
@ -15,6 +13,7 @@ FLOODSUB_PROTOCOL_TEST_CASES = [
|
||||
{
|
||||
"name": "simple_two_nodes",
|
||||
"supported_protocols": SUPPORTED_PROTOCOLS,
|
||||
"nodes": ["A", "B"],
|
||||
"adj_list": {"A": ["B"]},
|
||||
"topic_map": {"topic1": ["B"]},
|
||||
"messages": [{"topics": ["topic1"], "data": b"foo", "node_id": "A"}],
|
||||
@ -22,6 +21,7 @@ FLOODSUB_PROTOCOL_TEST_CASES = [
|
||||
{
|
||||
"name": "three_nodes_two_topics",
|
||||
"supported_protocols": SUPPORTED_PROTOCOLS,
|
||||
"nodes": ["A", "B", "C"],
|
||||
"adj_list": {"A": ["B"], "B": ["C"]},
|
||||
"topic_map": {"topic1": ["B", "C"], "topic2": ["B", "C"]},
|
||||
"messages": [
|
||||
@ -32,6 +32,7 @@ FLOODSUB_PROTOCOL_TEST_CASES = [
|
||||
{
|
||||
"name": "two_nodes_one_topic_single_subscriber_is_sender",
|
||||
"supported_protocols": SUPPORTED_PROTOCOLS,
|
||||
"nodes": ["A", "B"],
|
||||
"adj_list": {"A": ["B"]},
|
||||
"topic_map": {"topic1": ["B"]},
|
||||
"messages": [{"topics": ["topic1"], "data": b"Alex is tall", "node_id": "B"}],
|
||||
@ -39,6 +40,7 @@ FLOODSUB_PROTOCOL_TEST_CASES = [
|
||||
{
|
||||
"name": "two_nodes_one_topic_two_msgs",
|
||||
"supported_protocols": SUPPORTED_PROTOCOLS,
|
||||
"nodes": ["A", "B"],
|
||||
"adj_list": {"A": ["B"]},
|
||||
"topic_map": {"topic1": ["B"]},
|
||||
"messages": [
|
||||
@ -49,6 +51,7 @@ FLOODSUB_PROTOCOL_TEST_CASES = [
|
||||
{
|
||||
"name": "seven_nodes_tree_one_topics",
|
||||
"supported_protocols": SUPPORTED_PROTOCOLS,
|
||||
"nodes": ["1", "2", "3", "4", "5", "6", "7"],
|
||||
"adj_list": {"1": ["2", "3"], "2": ["4", "5"], "3": ["6", "7"]},
|
||||
"topic_map": {"astrophysics": ["2", "3", "4", "5", "6", "7"]},
|
||||
"messages": [{"topics": ["astrophysics"], "data": b"e=mc^2", "node_id": "1"}],
|
||||
@ -56,6 +59,7 @@ FLOODSUB_PROTOCOL_TEST_CASES = [
|
||||
{
|
||||
"name": "seven_nodes_tree_three_topics",
|
||||
"supported_protocols": SUPPORTED_PROTOCOLS,
|
||||
"nodes": ["1", "2", "3", "4", "5", "6", "7"],
|
||||
"adj_list": {"1": ["2", "3"], "2": ["4", "5"], "3": ["6", "7"]},
|
||||
"topic_map": {
|
||||
"astrophysics": ["2", "3", "4", "5", "6", "7"],
|
||||
@ -71,6 +75,7 @@ FLOODSUB_PROTOCOL_TEST_CASES = [
|
||||
{
|
||||
"name": "seven_nodes_tree_three_topics_diff_origin",
|
||||
"supported_protocols": SUPPORTED_PROTOCOLS,
|
||||
"nodes": ["1", "2", "3", "4", "5", "6", "7"],
|
||||
"adj_list": {"1": ["2", "3"], "2": ["4", "5"], "3": ["6", "7"]},
|
||||
"topic_map": {
|
||||
"astrophysics": ["1", "2", "3", "4", "5", "6", "7"],
|
||||
@ -86,6 +91,7 @@ FLOODSUB_PROTOCOL_TEST_CASES = [
|
||||
{
|
||||
"name": "three_nodes_clique_two_topic_diff_origin",
|
||||
"supported_protocols": SUPPORTED_PROTOCOLS,
|
||||
"nodes": ["1", "2", "3"],
|
||||
"adj_list": {"1": ["2", "3"], "2": ["3"]},
|
||||
"topic_map": {"astrophysics": ["1", "2", "3"], "school": ["1", "2", "3"]},
|
||||
"messages": [
|
||||
@ -97,6 +103,7 @@ FLOODSUB_PROTOCOL_TEST_CASES = [
|
||||
{
|
||||
"name": "four_nodes_clique_two_topic_diff_origin_many_msgs",
|
||||
"supported_protocols": SUPPORTED_PROTOCOLS,
|
||||
"nodes": ["1", "2", "3", "4"],
|
||||
"adj_list": {
|
||||
"1": ["2", "3", "4"],
|
||||
"2": ["1", "3", "4"],
|
||||
@ -120,6 +127,7 @@ FLOODSUB_PROTOCOL_TEST_CASES = [
|
||||
{
|
||||
"name": "five_nodes_ring_two_topic_diff_origin_many_msgs",
|
||||
"supported_protocols": SUPPORTED_PROTOCOLS,
|
||||
"nodes": ["1", "2", "3", "4", "5"],
|
||||
"adj_list": {"1": ["2"], "2": ["3"], "3": ["4"], "4": ["5"], "5": ["1"]},
|
||||
"topic_map": {
|
||||
"astrophysics": ["1", "2", "3", "4", "5"],
|
||||
@ -143,7 +151,7 @@ floodsub_protocol_pytest_params = [
|
||||
]
|
||||
|
||||
|
||||
async def perform_test_from_obj(obj, router_factory) -> None:
|
||||
async def perform_test_from_obj(obj, pubsub_factory) -> None:
|
||||
"""
|
||||
Perform pubsub tests from a test obj.
|
||||
test obj are composed as follows:
|
||||
@ -174,88 +182,75 @@ async def perform_test_from_obj(obj, router_factory) -> None:
|
||||
|
||||
# Step 1) Create graph
|
||||
adj_list = obj["adj_list"]
|
||||
node_list = obj["nodes"]
|
||||
node_map = {}
|
||||
pubsub_map = {}
|
||||
|
||||
async def add_node(node_id_str: str) -> None:
|
||||
pubsub_router = router_factory(protocols=obj["supported_protocols"])
|
||||
pubsub = PubsubFactory(router=pubsub_router)
|
||||
await pubsub.host.get_network().listen(LISTEN_MADDR)
|
||||
node_map[node_id_str] = pubsub.host
|
||||
pubsub_map[node_id_str] = pubsub
|
||||
async with pubsub_factory(
|
||||
number=len(node_list), protocols=obj["supported_protocols"]
|
||||
) as pubsubs:
|
||||
for node_id_str, pubsub in zip(node_list, pubsubs):
|
||||
node_map[node_id_str] = pubsub.host
|
||||
pubsub_map[node_id_str] = pubsub
|
||||
|
||||
tasks_connect = []
|
||||
for start_node_id in adj_list:
|
||||
# Create node if node does not yet exist
|
||||
if start_node_id not in node_map:
|
||||
await add_node(start_node_id)
|
||||
# Connect nodes and wait at least for 2 seconds
|
||||
async with trio.open_nursery() as nursery:
|
||||
for start_node_id in adj_list:
|
||||
# For each neighbor of start_node, create if does not yet exist,
|
||||
# then connect start_node to neighbor
|
||||
for neighbor_id in adj_list[start_node_id]:
|
||||
nursery.start_soon(
|
||||
connect, node_map[start_node_id], node_map[neighbor_id]
|
||||
)
|
||||
nursery.start_soon(trio.sleep, 2)
|
||||
|
||||
# For each neighbor of start_node, create if does not yet exist,
|
||||
# then connect start_node to neighbor
|
||||
for neighbor_id in adj_list[start_node_id]:
|
||||
# Create neighbor if neighbor does not yet exist
|
||||
if neighbor_id not in node_map:
|
||||
await add_node(neighbor_id)
|
||||
tasks_connect.append(
|
||||
connect(node_map[start_node_id], node_map[neighbor_id])
|
||||
)
|
||||
# Connect nodes and wait at least for 2 seconds
|
||||
await asyncio.gather(*tasks_connect, asyncio.sleep(2))
|
||||
# Step 2) Subscribe to topics
|
||||
queues_map = {}
|
||||
topic_map = obj["topic_map"]
|
||||
|
||||
# Step 2) Subscribe to topics
|
||||
queues_map = {}
|
||||
topic_map = obj["topic_map"]
|
||||
async def subscribe_node(node_id, topic):
|
||||
if node_id not in queues_map:
|
||||
queues_map[node_id] = {}
|
||||
# Avoid repeated works
|
||||
if topic in queues_map[node_id]:
|
||||
# Checkpoint
|
||||
await trio.hazmat.checkpoint()
|
||||
return
|
||||
sub = await pubsub_map[node_id].subscribe(topic)
|
||||
queues_map[node_id][topic] = sub
|
||||
|
||||
tasks_topic = []
|
||||
tasks_topic_data = []
|
||||
for topic, node_ids in topic_map.items():
|
||||
for node_id in node_ids:
|
||||
tasks_topic.append(pubsub_map[node_id].subscribe(topic))
|
||||
tasks_topic_data.append((node_id, topic))
|
||||
tasks_topic.append(asyncio.sleep(2))
|
||||
async with trio.open_nursery() as nursery:
|
||||
for topic, node_ids in topic_map.items():
|
||||
for node_id in node_ids:
|
||||
nursery.start_soon(subscribe_node, node_id, topic)
|
||||
nursery.start_soon(trio.sleep, 2)
|
||||
|
||||
# Gather is like Promise.all
|
||||
responses = await asyncio.gather(*tasks_topic)
|
||||
for i in range(len(responses) - 1):
|
||||
node_id, topic = tasks_topic_data[i]
|
||||
if node_id not in queues_map:
|
||||
queues_map[node_id] = {}
|
||||
# Store queue in topic-queue map for node
|
||||
queues_map[node_id][topic] = responses[i]
|
||||
# Step 3) Publish messages
|
||||
topics_in_msgs_ordered = []
|
||||
messages = obj["messages"]
|
||||
|
||||
# Allow time for subscribing before continuing
|
||||
await asyncio.sleep(0.01)
|
||||
for msg in messages:
|
||||
topics = msg["topics"]
|
||||
data = msg["data"]
|
||||
node_id = msg["node_id"]
|
||||
|
||||
# Step 3) Publish messages
|
||||
topics_in_msgs_ordered = []
|
||||
messages = obj["messages"]
|
||||
tasks_publish = []
|
||||
# Publish message
|
||||
# TODO: Should be single RPC package with several topics
|
||||
for topic in topics:
|
||||
await pubsub_map[node_id].publish(topic, data)
|
||||
|
||||
for msg in messages:
|
||||
topics = msg["topics"]
|
||||
data = msg["data"]
|
||||
node_id = msg["node_id"]
|
||||
# For each topic in topics, add (topic, node_id, data) tuple to ordered test list
|
||||
for topic in topics:
|
||||
topics_in_msgs_ordered.append((topic, node_id, data))
|
||||
# Allow time for publishing before continuing
|
||||
await trio.sleep(1)
|
||||
|
||||
# Publish message
|
||||
# TODO: Should be single RPC package with several topics
|
||||
for topic in topics:
|
||||
tasks_publish.append(pubsub_map[node_id].publish(topic, data))
|
||||
|
||||
# For each topic in topics, add (topic, node_id, data) tuple to ordered test list
|
||||
for topic in topics:
|
||||
topics_in_msgs_ordered.append((topic, node_id, data))
|
||||
|
||||
# Allow time for publishing before continuing
|
||||
await asyncio.gather(*tasks_publish, asyncio.sleep(2))
|
||||
|
||||
# Step 4) Check that all messages were received correctly.
|
||||
for topic, origin_node_id, data in topics_in_msgs_ordered:
|
||||
# Look at each node in each topic
|
||||
for node_id in topic_map[topic]:
|
||||
# Get message from subscription queue
|
||||
msg = await queues_map[node_id][topic].get()
|
||||
assert data == msg.data
|
||||
# Check the message origin
|
||||
assert node_map[origin_node_id].get_id().to_bytes() == msg.from_id
|
||||
|
||||
# Success, terminate pending tasks.
|
||||
# Step 4) Check that all messages were received correctly.
|
||||
for topic, origin_node_id, data in topics_in_msgs_ordered:
|
||||
# Look at each node in each topic
|
||||
for node_id in topic_map[topic]:
|
||||
# Get message from subscription queue
|
||||
msg = await queues_map[node_id][topic].receive()
|
||||
assert data == msg.data
|
||||
# Check the message origin
|
||||
assert node_map[origin_node_id].get_id().to_bytes() == msg.from_id
|
||||
|
||||
@ -1,17 +1,9 @@
|
||||
from typing import Callable, List, Sequence, Tuple
|
||||
from typing import Awaitable, Callable
|
||||
|
||||
import multiaddr
|
||||
import trio
|
||||
|
||||
from libp2p import new_node
|
||||
from libp2p.host.basic_host import BasicHost
|
||||
from libp2p.host.host_interface import IHost
|
||||
from libp2p.kademlia.network import KademliaServer
|
||||
from libp2p.network.stream.net_stream_interface import INetStream
|
||||
from libp2p.network.swarm import Swarm
|
||||
from libp2p.peer.peerinfo import info_from_p2p_addr
|
||||
from libp2p.routing.interfaces import IPeerRouting
|
||||
from libp2p.routing.kademlia.kademlia_peer_router import KadmeliaPeerRouter
|
||||
|
||||
from .constants import MAX_READ_LEN
|
||||
|
||||
@ -36,49 +28,9 @@ async def connect(node1: IHost, node2: IHost) -> None:
|
||||
await node1.connect(info)
|
||||
|
||||
|
||||
async def set_up_nodes_by_transport_opt(
|
||||
transport_opt_list: Sequence[Sequence[str]], nursery: trio.Nursery
|
||||
) -> Tuple[BasicHost, ...]:
|
||||
nodes_list = []
|
||||
for transport_opt in transport_opt_list:
|
||||
node = new_node(transport_opt=transport_opt)
|
||||
await node.get_network().listen(
|
||||
multiaddr.Multiaddr(transport_opt[0]), nursery=nursery
|
||||
)
|
||||
nodes_list.append(node)
|
||||
return tuple(nodes_list)
|
||||
|
||||
|
||||
async def set_up_nodes_by_transport_and_disc_opt(
|
||||
transport_disc_opt_list: Sequence[Tuple[Sequence[str], IPeerRouting]]
|
||||
) -> Tuple[BasicHost, ...]:
|
||||
nodes_list = []
|
||||
for transport_opt, disc_opt in transport_disc_opt_list:
|
||||
node = await new_node(transport_opt=transport_opt, disc_opt=disc_opt)
|
||||
await node.get_network().listen(multiaddr.Multiaddr(transport_opt[0]))
|
||||
nodes_list.append(node)
|
||||
return tuple(nodes_list)
|
||||
|
||||
|
||||
async def set_up_routers(
|
||||
router_ports: Tuple[int, ...] = (0, 0)
|
||||
) -> List[KadmeliaPeerRouter]:
|
||||
"""The default ``router_confs`` selects two free ports local to this
|
||||
machine."""
|
||||
bootstrap_node = KademliaServer() # type: ignore
|
||||
await bootstrap_node.listen(router_ports[0])
|
||||
|
||||
routers = [KadmeliaPeerRouter(bootstrap_node)]
|
||||
for port in router_ports[1:]:
|
||||
node = KademliaServer() # type: ignore
|
||||
await node.listen(port)
|
||||
|
||||
await node.bootstrap_node(bootstrap_node.address)
|
||||
routers.append(KadmeliaPeerRouter(node))
|
||||
return routers
|
||||
|
||||
|
||||
def create_echo_stream_handler(ack_prefix: str) -> Callable[[INetStream], None]:
|
||||
def create_echo_stream_handler(
|
||||
ack_prefix: str
|
||||
) -> Callable[[INetStream], Awaitable[None]]:
|
||||
async def echo_stream_handler(stream: INetStream) -> None:
|
||||
while True:
|
||||
read_string = (await stream.read(MAX_READ_LEN)).decode()
|
||||
|
||||
Reference in New Issue
Block a user