chore: refactor flaky test with while (#1698)

waku-org · Apr 25, 2023 · dca0e9b · dca0e9b
1 parent e041e04
commit dca0e9b
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 17 deletions.
diff --git 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)
-    check:
-      connOpt.isSome
-
-    # Give the algorithm some time to work its magic
-    await sleepAsync(2000.millis)
+    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
 
-    ## 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
@@ -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.