From 99eabe49eb5e5123ee7e29f576be7918e79d81b4 Mon Sep 17 00:00:00 2001
From: NIC619 <twedusuck@gmail.com>
Date: Sun, 21 Jul 2019 22:28:43 +0800
Subject: [PATCH] Add `handle_prune` test

---
 tests/pubsub/test_gossipsub.py | 46 ++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

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