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 => { + 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);