From b405fd76e95dee070c556f3921bb2473fe1b765f Mon Sep 17 00:00:00 2001
From: NIC619 <twedusuck@gmail.com>
Date: Tue, 3 Dec 2019 15:49:58 +0800
Subject: [PATCH] Add test for gossip heartbeat

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

diff --git a/tests/pubsub/test_gossipsub.py b/tests/pubsub/test_gossipsub.py
index 3ff602d1..a8a89a6c 100644
--- a/tests/pubsub/test_gossipsub.py
+++ b/tests/pubsub/test_gossipsub.py
@@ -406,3 +406,63 @@ async def test_mesh_heartbeat(
             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
+):
+    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
+    peer_topics = {
+        topic_mesh: fake_peer_ids[:topic_mesh_peer_count],
+        topic_fanout: fake_peer_ids[topic_mesh_peer_count:],
+    }
+    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.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.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.setattr(pubsubs_gsub[0].router.mcache, "window", window)
+
+    peers_to_gossip = pubsubs_gsub[0].router.gossip_heartbeat()
+    if topic_mesh_peer_count - initial_peer_count < GOSSIPSUB_PARAMS.degree:
+        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]:
+            assert peer not in mesh_peers
+            assert topic_mesh in peers_to_gossip[peer]
+        elif peer in peer_topics[topic_fanout]:
+            assert peer not in fanout_peers
+            assert topic_fanout in peers_to_gossip[peer]