From eb1e8d49074839dd2dce75bf0ebf6703c098836d Mon Sep 17 00:00:00 2001
From: HDegroote <75906619+HDegroote@users.noreply.github.com>
Date: Fri, 27 Sep 2024 20:44:46 +0200
Subject: [PATCH] Make test more robust

---
 test/swarm.js | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/test/swarm.js b/test/swarm.js
index f148092..c548497 100644
--- a/test/swarm.js
+++ b/test/swarm.js
@@ -616,9 +616,15 @@ test('peer-discovery object deleted when corresponding connection closes (server
     connected.pass('swarm2')
     conn.on('error', noop)
   })
+
+  let resolveConnClosed = null
+  const connClosed = new Promise(resolve => {
+    resolveConnClosed = resolve
+  })
   swarm1.on('connection', (conn) => {
     otherConnected.pass('swarm1')
     conn.on('error', noop)
+    conn.on('close', resolveConnClosed)
   })
 
   const topic = Buffer.alloc(32).fill('hello world')
@@ -634,7 +640,7 @@ test('peer-discovery object deleted when corresponding connection closes (server
   await swarm2.destroy()
 
   // Ensure other side detects closed connection
-  await eventFlush()
+  await connClosed
 
   t.is(swarm1.peers.size, 0, 'No peerInfo memory leak')
 
@@ -791,7 +797,3 @@ test('topic and peer get unslabbed in PeerInfo', async (t) => {
 })
 
 function noop () {}
-
-function eventFlush () {
-  return new Promise(resolve => setImmediate(resolve))
-}