import asyncio import random import pytest from libp2p.peer.id import ID from libp2p.tools.constants import GOSSIPSUB_PARAMS, GossipsubParams from libp2p.tools.pubsub.utils import dense_connect, one_to_all_connect from libp2p.tools.utils import connect @pytest.mark.parametrize( "num_hosts, gossipsub_params", ((4, GossipsubParams(degree=4, degree_low=3, degree_high=5)),), ) @pytest.mark.asyncio async def test_join(num_hosts, hosts, pubsubs_gsub): gossipsubs = tuple(pubsub.router for pubsub in pubsubs_gsub) hosts_indices = list(range(num_hosts)) topic = "test_join" central_node_index = 0 # Remove index of central host from the indices hosts_indices.remove(central_node_index) num_subscribed_peer = 2 subscribed_peer_indices = random.sample(hosts_indices, num_subscribed_peer) # All pubsub except the one of central node subscribe to topic for i in subscribed_peer_indices: await pubsubs_gsub[i].subscribe(topic) # Connect central host to all other hosts await one_to_all_connect(hosts, central_node_index) # Wait 2 seconds for heartbeat to allow mesh to connect await asyncio.sleep(2) # Central node publish to the topic so that this topic # is added to central node's fanout # publish from the randomly chosen host await pubsubs_gsub[central_node_index].publish(topic, b"data") # Check that the gossipsub of central node has fanout for the topic assert topic in gossipsubs[central_node_index].fanout # Check that the gossipsub of central node does not have a mesh for the topic assert topic not in gossipsubs[central_node_index].mesh # Central node subscribes the topic await pubsubs_gsub[central_node_index].subscribe(topic) await asyncio.sleep(2) # Check that the gossipsub of central node no longer has fanout for the topic assert topic not in gossipsubs[central_node_index].fanout for i in hosts_indices: if i in subscribed_peer_indices: assert hosts[i].get_id() in gossipsubs[central_node_index].mesh[topic] assert hosts[central_node_index].get_id() in gossipsubs[i].mesh[topic] else: assert hosts[i].get_id() not in gossipsubs[central_node_index].mesh[topic] assert topic not in gossipsubs[i].mesh @pytest.mark.parametrize("num_hosts", (1,)) @pytest.mark.asyncio async def test_leave(pubsubs_gsub): gossipsub = pubsubs_gsub[0].router topic = "test_leave" assert topic not in gossipsub.mesh await gossipsub.join(topic) assert topic in gossipsub.mesh await gossipsub.leave(topic) assert topic not in gossipsub.mesh # Test re-leave await gossipsub.leave(topic) @pytest.mark.parametrize("num_hosts", (2,)) @pytest.mark.asyncio async def test_handle_graft(pubsubs_gsub, hosts, event_loop, monkeypatch): gossipsubs = tuple(pubsub.router for pubsub in pubsubs_gsub) index_alice = 0 id_alice = hosts[index_alice].get_id() index_bob = 1 id_bob = hosts[index_bob].get_id() await connect(hosts[index_alice], hosts[index_bob]) # Wait 2 seconds for heartbeat to allow mesh to connect await asyncio.sleep(2) topic = "test_handle_graft" # Only lice subscribe to the topic await gossipsubs[index_alice].join(topic) # Monkey patch bob's `emit_prune` function so we can # check if it is called in `handle_graft` event_emit_prune = asyncio.Event() async def emit_prune(topic, sender_peer_id): event_emit_prune.set() monkeypatch.setattr(gossipsubs[index_bob], "emit_prune", emit_prune) # Check that alice is bob's peer but not his mesh peer assert id_alice in gossipsubs[index_bob].peers_gossipsub assert topic not in gossipsubs[index_bob].mesh await gossipsubs[index_alice].emit_graft(topic, id_bob) # Check that `emit_prune` is called await asyncio.wait_for(event_emit_prune.wait(), timeout=1, loop=event_loop) assert event_emit_prune.is_set() # Check that bob is alice's peer but not her mesh peer assert topic in gossipsubs[index_alice].mesh assert id_bob not in gossipsubs[index_alice].mesh[topic] assert id_bob in gossipsubs[index_alice].peers_gossipsub await gossipsubs[index_bob].emit_graft(topic, id_alice) await asyncio.sleep(1) # Check that bob is now alice's mesh peer assert id_bob in gossipsubs[index_alice].mesh[topic] @pytest.mark.parametrize( "num_hosts, gossipsub_params", ((2, GossipsubParams(heartbeat_interval=3)),) ) @pytest.mark.asyncio async def test_handle_prune(pubsubs_gsub, hosts): gossipsubs = tuple(pubsub.router for pubsub in pubsubs_gsub) index_alice = 0 id_alice = hosts[index_alice].get_id() index_bob = 1 id_bob = hosts[index_bob].get_id() topic = "test_handle_prune" for pubsub in pubsubs_gsub: await pubsub.subscribe(topic) await connect(hosts[index_alice], hosts[index_bob]) # Wait for heartbeat to allow mesh to connect await asyncio.sleep(1) # Check that they are each other's mesh peer assert id_alice in gossipsubs[index_bob].mesh[topic] assert id_bob in gossipsubs[index_alice].mesh[topic] # alice emit prune message to bob, alice should be removed # from bob's mesh peer await gossipsubs[index_alice].emit_prune(topic, id_bob) # `emit_prune` does not remove bob from alice's mesh peers assert id_bob in gossipsubs[index_alice].mesh[topic] # NOTE: We increase `heartbeat_interval` to 3 seconds so that bob will not # add alice back to his mesh after heartbeat. # Wait for bob to `handle_prune` await asyncio.sleep(0.1) # Check that alice is no longer bob's mesh peer assert id_alice not in gossipsubs[index_bob].mesh[topic] @pytest.mark.parametrize("num_hosts", (10,)) @pytest.mark.asyncio async def test_dense(num_hosts, pubsubs_gsub, hosts): num_msgs = 5 # All pubsub subscribe to foobar queues = [] for pubsub in pubsubs_gsub: q = await pubsub.subscribe("foobar") # Add each blocking queue to an array of blocking queues queues.append(q) # Densely connect libp2p hosts in a random way await dense_connect(hosts) # Wait 2 seconds for heartbeat to allow mesh to connect await asyncio.sleep(2) for i in range(num_msgs): msg_content = b"foo " + i.to_bytes(1, "big") # randomly pick a message origin origin_idx = random.randint(0, num_hosts - 1) # publish from the randomly chosen host await pubsubs_gsub[origin_idx].publish("foobar", msg_content) await asyncio.sleep(0.5) # Assert that all blocking queues receive the message for queue in queues: msg = await queue.get() assert msg.data == msg_content @pytest.mark.parametrize("num_hosts", (10,)) @pytest.mark.asyncio async def test_fanout(hosts, pubsubs_gsub): num_msgs = 5 # All pubsub subscribe to foobar except for `pubsubs_gsub[0]` queues = [] for i in range(1, len(pubsubs_gsub)): q = await pubsubs_gsub[i].subscribe("foobar") # Add each blocking queue to an array of blocking queues queues.append(q) # Sparsely connect libp2p hosts in random way await dense_connect(hosts) # Wait 2 seconds for heartbeat to allow mesh to connect await asyncio.sleep(2) topic = "foobar" # Send messages with origin not subscribed for i in range(num_msgs): msg_content = b"foo " + i.to_bytes(1, "big") # Pick the message origin to the node that is not subscribed to 'foobar' origin_idx = 0 # publish from the randomly chosen host await pubsubs_gsub[origin_idx].publish(topic, msg_content) await asyncio.sleep(0.5) # Assert that all blocking queues receive the message for queue in queues: msg = await queue.get() assert msg.data == msg_content # Subscribe message origin queues.insert(0, await pubsubs_gsub[0].subscribe(topic)) # Send messages again for i in range(num_msgs): msg_content = b"bar " + i.to_bytes(1, "big") # Pick the message origin to the node that is not subscribed to 'foobar' origin_idx = 0 # publish from the randomly chosen host await pubsubs_gsub[origin_idx].publish(topic, msg_content) await asyncio.sleep(0.5) # Assert that all blocking queues receive the message for queue in queues: msg = await queue.get() assert msg.data == msg_content @pytest.mark.parametrize("num_hosts", (10,)) @pytest.mark.asyncio @pytest.mark.slow async def test_fanout_maintenance(hosts, pubsubs_gsub): num_msgs = 5 # All pubsub subscribe to foobar queues = [] topic = "foobar" for i in range(1, len(pubsubs_gsub)): q = await pubsubs_gsub[i].subscribe(topic) # Add each blocking queue to an array of blocking queues queues.append(q) # Sparsely connect libp2p hosts in random way await dense_connect(hosts) # Wait 2 seconds for heartbeat to allow mesh to connect await asyncio.sleep(2) # Send messages with origin not subscribed for i in range(num_msgs): msg_content = b"foo " + i.to_bytes(1, "big") # Pick the message origin to the node that is not subscribed to 'foobar' origin_idx = 0 # publish from the randomly chosen host await pubsubs_gsub[origin_idx].publish(topic, msg_content) await asyncio.sleep(0.5) # Assert that all blocking queues receive the message for queue in queues: msg = await queue.get() assert msg.data == msg_content for sub in pubsubs_gsub: await sub.unsubscribe(topic) queues = [] await asyncio.sleep(2) # Resub and repeat for i in range(1, len(pubsubs_gsub)): q = await pubsubs_gsub[i].subscribe(topic) # Add each blocking queue to an array of blocking queues queues.append(q) await asyncio.sleep(2) # Check messages can still be sent for i in range(num_msgs): msg_content = b"bar " + i.to_bytes(1, "big") # Pick the message origin to the node that is not subscribed to 'foobar' origin_idx = 0 # publish from the randomly chosen host await pubsubs_gsub[origin_idx].publish(topic, msg_content) await asyncio.sleep(0.5) # Assert that all blocking queues receive the message for queue in queues: msg = await queue.get() assert msg.data == msg_content @pytest.mark.parametrize( "num_hosts, gossipsub_params", ( ( 2, GossipsubParams( degree=1, degree_low=0, degree_high=2, gossip_window=50, gossip_history=100, ), ), ), ) @pytest.mark.asyncio async def test_gossip_propagation(hosts, pubsubs_gsub): topic = "foo" await pubsubs_gsub[0].subscribe(topic) # node 0 publish to topic msg_content = b"foo_msg" # publish from the randomly chosen host await pubsubs_gsub[0].publish(topic, msg_content) # now node 1 subscribes queue_1 = await pubsubs_gsub[1].subscribe(topic) await connect(hosts[0], hosts[1]) # wait for gossip heartbeat await asyncio.sleep(2) # should be able to read message msg = await queue_1.get() assert msg.data == msg_content @pytest.mark.parametrize( "num_hosts, gossipsub_params", ((1, GossipsubParams(heartbeat_initial_delay=100)),) ) @pytest.mark.parametrize("initial_mesh_peer_count", (7, 10, 13)) @pytest.mark.asyncio async def test_mesh_heartbeat( num_hosts, initial_mesh_peer_count, pubsubs_gsub, hosts, monkeypatch ): # It's difficult to set up the initial peer subscription condition. # Ideally I would like to have initial mesh peer count that's below ``GossipSubDegree`` # so I can test if `mesh_heartbeat` return correct peers to GRAFT. # The problem is that I can not set it up so that we have peers subscribe to the topic # but not being part of our mesh peers (as these peers are the peers to GRAFT). # So I monkeypatch the peer subscriptions and our mesh peers. total_peer_count = 14 topic = "TEST_MESH_HEARTBEAT" fake_peer_ids = [ ID((i).to_bytes(2, byteorder="big")) for i in range(total_peer_count) ] monkeypatch.setattr(pubsubs_gsub[0].router, "peers_gossipsub", fake_peer_ids) peer_topics = {topic: fake_peer_ids} # Monkeypatch the peer subscriptions monkeypatch.setattr(pubsubs_gsub[0], "peer_topics", peer_topics) mesh_peer_indices = random.sample(range(total_peer_count), initial_mesh_peer_count) mesh_peers = [fake_peer_ids[i] for i in mesh_peer_indices] router_mesh = {topic: list(mesh_peers)} # Monkeypatch our mesh peers monkeypatch.setattr(pubsubs_gsub[0].router, "mesh", router_mesh) peers_to_graft, peers_to_prune = pubsubs_gsub[0].router.mesh_heartbeat() if initial_mesh_peer_count > GOSSIPSUB_PARAMS.degree: # If number of initial mesh peers is more than `GossipSubDegree`, we should PRUNE mesh peers assert len(peers_to_graft) == 0 assert len(peers_to_prune) == initial_mesh_peer_count - GOSSIPSUB_PARAMS.degree for peer in peers_to_prune: assert peer in mesh_peers elif initial_mesh_peer_count < GOSSIPSUB_PARAMS.degree: # If number of initial mesh peers is less than `GossipSubDegree`, we should GRAFT more peers assert len(peers_to_prune) == 0 assert len(peers_to_graft) == GOSSIPSUB_PARAMS.degree - initial_mesh_peer_count for peer in peers_to_graft: assert peer not in mesh_peers else: assert len(peers_to_prune) == 0 and len(peers_to_graft) == 0 @pytest.mark.parametrize( "num_hosts, gossipsub_params", ((1, GossipsubParams(heartbeat_initial_delay=100)),) ) @pytest.mark.parametrize("initial_peer_count", (1, 4, 7)) @pytest.mark.asyncio async def test_gossip_heartbeat( num_hosts, initial_peer_count, pubsubs_gsub, hosts, monkeypatch ): # The problem is that I can not set it up so that we have peers subscribe to the topic # but not being part of our mesh peers (as these peers are the peers to GRAFT). # So I monkeypatch the peer subscriptions and our mesh peers. total_peer_count = 28 topic_mesh = "TEST_GOSSIP_HEARTBEAT_1" topic_fanout = "TEST_GOSSIP_HEARTBEAT_2" fake_peer_ids = [ ID((i).to_bytes(2, byteorder="big")) for i in range(total_peer_count) ] monkeypatch.setattr(pubsubs_gsub[0].router, "peers_gossipsub", fake_peer_ids) topic_mesh_peer_count = 14 # Split into mesh peers and fanout peers peer_topics = { topic_mesh: fake_peer_ids[:topic_mesh_peer_count], topic_fanout: fake_peer_ids[topic_mesh_peer_count:], } # Monkeypatch the peer subscriptions monkeypatch.setattr(pubsubs_gsub[0], "peer_topics", peer_topics) mesh_peer_indices = random.sample(range(topic_mesh_peer_count), initial_peer_count) mesh_peers = [fake_peer_ids[i] for i in mesh_peer_indices] router_mesh = {topic_mesh: list(mesh_peers)} # Monkeypatch our mesh peers monkeypatch.setattr(pubsubs_gsub[0].router, "mesh", router_mesh) fanout_peer_indices = random.sample( range(topic_mesh_peer_count, total_peer_count), initial_peer_count ) fanout_peers = [fake_peer_ids[i] for i in fanout_peer_indices] router_fanout = {topic_fanout: list(fanout_peers)} # Monkeypatch our fanout peers monkeypatch.setattr(pubsubs_gsub[0].router, "fanout", router_fanout) def window(topic): if topic == topic_mesh: return [topic_mesh] elif topic == topic_fanout: return [topic_fanout] else: return [] # Monkeypatch the memory cache messages monkeypatch.setattr(pubsubs_gsub[0].router.mcache, "window", window) peers_to_gossip = pubsubs_gsub[0].router.gossip_heartbeat() # If our mesh peer count is less than `GossipSubDegree`, we should gossip to up to # `GossipSubDegree` peers (exclude mesh peers). if topic_mesh_peer_count - initial_peer_count < GOSSIPSUB_PARAMS.degree: # The same goes for fanout so it's two times the number of peers to gossip. assert len(peers_to_gossip) == 2 * (topic_mesh_peer_count - initial_peer_count) elif topic_mesh_peer_count - initial_peer_count >= GOSSIPSUB_PARAMS.degree: assert len(peers_to_gossip) == 2 * (GOSSIPSUB_PARAMS.degree) for peer in peers_to_gossip: if peer in peer_topics[topic_mesh]: # Check that the peer to gossip to is not in our mesh peers assert peer not in mesh_peers assert topic_mesh in peers_to_gossip[peer] elif peer in peer_topics[topic_fanout]: # Check that the peer to gossip to is not in our fanout peers assert peer not in fanout_peers assert topic_fanout in peers_to_gossip[peer]