From b32badfe9a6af88c16e839c7fe17f0624152561f Mon Sep 17 00:00:00 2001
From: Alejandro Cabeza Romero <alex93cabeza@gmail.com>
Date: Mon, 26 Aug 2024 19:36:17 +0200
Subject: [PATCH] Add active waitForMesh proc.

---
 tests/pubsub/utils.nim | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/tests/pubsub/utils.nim b/tests/pubsub/utils.nim
index 4eb6a3821..937951954 100644
--- a/tests/pubsub/utils.nim
+++ b/tests/pubsub/utils.nim
@@ -147,6 +147,12 @@ proc subscribeRandom*(nodes: seq[PubSub]) {.async.} =
           await dialer.switch.connect(node.peerInfo.peerId, node.peerInfo.addrs)
           dialed.add(node.peerInfo.peerId)
 
+proc activeWait(
+    interval: Duration, maximum: Moment, timeoutErrorMessage = "Timeout on activeWait"
+) {.async.} =
+  await sleepAsync(interval)
+  doAssert Moment.now() < maximum, timeoutErrorMessage
+
 proc waitSub*(sender, receiver: auto, key: string) {.async.} =
   if sender == receiver:
     return
@@ -201,3 +207,19 @@ proc waitSubGraph*(nodes: seq[PubSub], key: string) {.async.} =
 
     await sleepAsync(5.milliseconds)
     doAssert Moment.now() < timeout, "waitSubGraph timeout!"
+
+proc waitForMesh*(
+    sender: auto, receiver: auto, key: string, timeoutDuration = 5.seconds
+) {.async.} =
+  if sender == receiver:
+    return
+
+  let
+    timeoutMoment = Moment.now() + timeoutDuration
+    gossipsubSender = GossipSub(sender)
+    receiverPeerId = receiver.peerInfo.peerId
+
+  while not gossipsubSender.mesh.hasPeerId(key, receiverPeerId):
+    trace "waitForMesh sleeping..."
+    await activeWait(5.milliseconds, timeoutMoment, "waitForMesh timeout!")
+  echo "waitForMesh done!"