From dca0e9b2ccdb021584d5e7047970bee61729f9cb Mon Sep 17 00:00:00 2001
From: Alvaro Revuelta <alvrevuelta@gmail.com>
Date: Tue, 25 Apr 2023 09:50:13 +0200
Subject: [PATCH] chore: refactor flaky test with while (#1698)

---
 tests/v2/test_waku_peer_exchange.nim        | 34 ++++++++++++---------
 tests/v2/waku_relay/test_wakunode_relay.nim | 13 ++++++--
 2 files changed, 29 insertions(+), 18 deletions(-)

diff --git a/tests/v2/test_waku_peer_exchange.nim b/tests/v2/test_waku_peer_exchange.nim
index ee22d1a0a..0e87239a9 100644
--- a/tests/v2/test_waku_peer_exchange.nim
+++ b/tests/v2/test_waku_peer_exchange.nim
@@ -126,8 +126,12 @@ procSuite "Waku Peer Exchange":
     await allFutures([node1.start(), node2.start(), node3.start()])
     await allFutures([node1.startDiscv5(), node2.startDiscv5()])
 
-    # Give disv5 some time to discover each other
-    await sleepAsync(5000.millis)
+    var attempts = 10
+    while (node1.wakuDiscv5.protocol.nodesDiscovered < 1 or
+          node2.wakuDiscv5.protocol.nodesDiscovered < 1) and
+          attempts > 0:
+      await sleepAsync(1.seconds)
+      attempts -= 1
 
     # node2 can be connected, so will be returned by peer exchange
     require (await node1.peerManager.connectRelay(node2.switch.peerInfo.toRemotePeerInfo()))
@@ -136,21 +140,21 @@ procSuite "Waku Peer Exchange":
     await node1.mountPeerExchange()
     await node3.mountPeerExchange()
 
-    let connOpt = await node3.peerManager.dialPeer(node1.switch.peerInfo.toRemotePeerInfo(), WakuPeerExchangeCodec)
+    var peerInfosLen = 0
+    var response: WakuPeerExchangeResult[PeerExchangeResponse]
+    attempts = 10
+    while peerInfosLen == 0 and attempts > 0:
+      var connOpt = await node3.peerManager.dialPeer(node1.switch.peerInfo.toRemotePeerInfo(), WakuPeerExchangeCodec)
+      require connOpt.isSome
+      response = await node3.wakuPeerExchange.request(1, connOpt.get())
+      require response.isOk
+      peerInfosLen = response.get().peerInfos.len
+      await sleepAsync(1.seconds)
+      attempts -= 1
+
     check:
-      connOpt.isSome
-
-    # Give the algorithm some time to work its magic
-    await sleepAsync(2000.millis)
-
-    ## When
-    let response = await node3.wakuPeerExchange.request(1, connOpt.get())
-
-    ## Then
-    check:
-      response.isOk
       response.get().peerInfos.len == 1
-      node1.wakuDiscv5.protocol.nodesDiscovered > 0
+      response.get().peerInfos[0].enr == node2.wakuDiscV5.protocol.localNode.record.raw
 
     await allFutures([node1.stop(), node2.stop(), node3.stop()])
 
diff --git a/tests/v2/waku_relay/test_wakunode_relay.nim b/tests/v2/waku_relay/test_wakunode_relay.nim
index 4c7cd3b37..95a4ad7f7 100644
--- a/tests/v2/waku_relay/test_wakunode_relay.nim
+++ b/tests/v2/waku_relay/test_wakunode_relay.nim
@@ -280,6 +280,9 @@ suite "WakuNode - Relay":
         let connOk = await nodes[i].peerManager.connectRelay(nodes[j].switch.peerInfo.toRemotePeerInfo())
         require connOk
 
+    # Connection triggers different actions, wait for them
+    await sleepAsync(500.millis)
+
     var msgReceived = 0
     proc handler(pubsubTopic: PubsubTopic, data: WakuMessage) {.async, gcsafe.} =
       msgReceived += 1
@@ -301,7 +304,7 @@ suite "WakuNode - Relay":
         await nodes[i].publish(spamProtectedTopic, msg)
 
     # Wait for gossip
-    await sleepAsync(1.seconds)
+    await sleepAsync(2.seconds)
 
     # 50 messages were sent to 5 peers = 250 messages
     check:
@@ -353,6 +356,9 @@ suite "WakuNode - Relay":
     proc handler(pubsubTopic: PubsubTopic, msg: WakuMessage) {.async, gcsafe.} =
       msgReceived += 1
 
+    # Connection triggers different actions, wait for them
+    await sleepAsync(500.millis)
+
     # Subscribe all nodes to the same topic/handler
     for node in nodes: node.wakuRelay.subscribe(spamProtectedTopic, handler)
     await sleepAsync(500.millis)
@@ -378,7 +384,7 @@ suite "WakuNode - Relay":
         await nodes[i].publish(spamProtectedTopic, unsignedMessage)
 
     # Wait for gossip
-    await sleepAsync(1.seconds)
+    await sleepAsync(2.seconds)
 
     # Since we have a full mesh with 5 nodes and each one publishes 50+50 msgs
     # there are 500 messages being sent.
@@ -442,6 +448,7 @@ suite "WakuNode - Relay":
         let connOk2 = await nodes[i].peerManager.connectRelay(nodes[j].switch.peerInfo.toRemotePeerInfo())
         require connOk2
 
+    # Connection triggers different actions, wait for them
     await sleepAsync(500.millis)
 
     # nodes[0] spams 50 non signed messages (nodes[0] just knows of nodes[1])
@@ -461,7 +468,7 @@ suite "WakuNode - Relay":
       await nodes[0].publish(spamProtectedTopic, msg)
 
     # Wait for gossip
-    await sleepAsync(1.seconds)
+    await sleepAsync(2.seconds)
 
     # only 100 messages are received (50 + 50) which demonstrate
     # nodes[1] doest gossip invalid messages.