diff --git a/tests/pubsub/test_gossipsub.py b/tests/pubsub/test_gossipsub.py index 9784b05..b827f9c 100644 --- a/tests/pubsub/test_gossipsub.py +++ b/tests/pubsub/test_gossipsub.py @@ -132,6 +132,7 @@ async def test_handle_graft(event_loop, monkeypatch): # 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 @@ -146,6 +147,7 @@ async def test_handle_graft(event_loop, monkeypatch): 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) @@ -156,6 +158,50 @@ async def test_handle_graft(event_loop, monkeypatch): await cleanup() +@pytest.mark.asyncio +async def test_handle_prune(): + num_hosts = 2 + libp2p_hosts = await create_libp2p_hosts(num_hosts) + + # Create pubsub, gossipsub instances + _, gossipsubs = create_pubsub_and_gossipsub_instances(libp2p_hosts, \ + SUPPORTED_PROTOCOLS, \ + 10, 9, 11, 30, 3, 5, 2) + + index_alice = 0 + id_alice = str(libp2p_hosts[index_alice].get_id()) + index_bob = 1 + id_bob = str(libp2p_hosts[index_bob].get_id()) + + topic = "test_handle_prune" + await gossipsubs[index_alice].join(topic) + await gossipsubs[index_bob].join(topic) + + await connect(libp2p_hosts[index_alice], libp2p_hosts[index_bob]) + + # Wait 2 seconds for heartbeat to allow mesh to connect + await asyncio.sleep(2) + + # 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) + + # FIXME: This test currently works because the heartbeat interval + # is increased to 2 seconds, so alice won't get add back into + # bob's mesh peer during heartbeat. + await asyncio.sleep(1) + + # Check that alice is no longer bob's mesh peer + assert id_alice not in gossipsubs[index_bob].mesh[topic] + assert id_bob in gossipsubs[index_alice].mesh[topic] + + await cleanup() + + @pytest.mark.asyncio async def test_dense(): # Create libp2p hosts