From 56b4de6fa3a9072eecf94a557d41277f7f09d357 Mon Sep 17 00:00:00 2001
From: spypsy <spypsy@users.noreply.github.com>
Date: Wed, 24 Apr 2024 15:03:49 +0100
Subject: [PATCH] fix: make discv5 test deterministic (#5968)

---
 .../p2p/src/service/discv5_service.test.ts    | 25 +++++++++++++++----
 1 file changed, 20 insertions(+), 5 deletions(-)

diff --git a/yarn-project/p2p/src/service/discv5_service.test.ts b/yarn-project/p2p/src/service/discv5_service.test.ts
index 9878379ea73..9228e1b87d0 100644
--- a/yarn-project/p2p/src/service/discv5_service.test.ts
+++ b/yarn-project/p2p/src/service/discv5_service.test.ts
@@ -1,11 +1,25 @@
-import { sleep } from '@aztec/foundation/sleep';
-
 import type { PeerId } from '@libp2p/interface';
 
 import { BootstrapNode } from '../bootstrap/bootstrap.js';
 import { DiscV5Service, PeerDiscoveryState } from './discV5_service.js';
 import { createLibP2PPeerId } from './libp2p_service.js';
 
+const waitForPeers = (node: DiscV5Service, expectedCount: number): Promise<void> => {
+  const timeout = 5_000;
+  return new Promise((resolve, reject) => {
+    const timeoutId = setTimeout(() => {
+      reject(new Error(`Timeout: Failed to connect to ${expectedCount} peers within ${timeout} ms`));
+    }, timeout);
+
+    node.on('peer:discovered', () => {
+      if (node.getAllPeers().length >= expectedCount) {
+        clearTimeout(timeoutId);
+        resolve();
+      }
+    });
+  });
+};
+
 describe('Discv5Service', () => {
   let bootNode: BootstrapNode;
   let bootNodePeerId: PeerId;
@@ -48,7 +62,8 @@ describe('Discv5Service', () => {
     const node2 = await createNode(port);
     await node1.start();
     await node2.start();
-    await sleep(200);
+    await waitForPeers(node2, 2);
+
     const node1Peers = await Promise.all(node1.getAllPeers().map(async peer => (await peer.peerId()).toString()));
     const node2Peers = await Promise.all(node2.getAllPeers().map(async peer => (await peer.peerId()).toString()));
 
@@ -68,13 +83,13 @@ describe('Discv5Service', () => {
     const node2 = await createNode(port);
     await node1.start();
     await node2.start();
-    await sleep(200);
+    await waitForPeers(node2, 2);
 
     await node2.stop();
     await bootNode.stop();
 
     await node2.start();
-    await sleep(200);
+    await waitForPeers(node2, 1);
 
     const node2Peers = await Promise.all(node2.getAllPeers().map(async peer => (await peer.peerId()).toString()));
     expect(node2Peers).toHaveLength(1);