optimize waitFor utility

waku-org · Jan 9, 2024 · dca2f6e · dca2f6e
1 parent 1ab3d58
commit dca2f6e
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 3 deletions.
diff --git a/packages/tests/src/waitForRemotePeerWithCodec.ts b/packages/tests/src/waitForRemotePeerWithCodec.ts
@@ -1,4 +1,5 @@
 import type { IdentifyResult } from "@libp2p/interface";
+import type { PeerId } from "@libp2p/interface/peer-id";
 import type { LightNode } from "@waku/interfaces";
 
 /**
@@ -12,8 +13,18 @@ import type { LightNode } from "@waku/interfaces";
  */
 export async function waitForRemotePeerWithCodec(
   waku: LightNode,
-  codec: string
+  codec: string,
+  nodePeerId: PeerId
 ): Promise<void> {
+  const connectedPeers = waku.libp2p
+    .getConnections()
+    .map((conn) => conn.remotePeer);
+  if (
+    connectedPeers.find((connectedPeer) => connectedPeer.equals(nodePeerId))
+  ) {
+    return;
+  }
+
   await new Promise<void>((resolve) => {
     const cb = (evt: CustomEvent<IdentifyResult>): void => {
       if (evt.detail.protocols.includes(codec)) {

diff --git a/packages/tests/tests/peer_exchange.node.spec.ts b/packages/tests/tests/peer_exchange.node.spec.ts
@@ -1,3 +1,4 @@
+import type { PeerId } from "@libp2p/interface/peer-id";
 import tests from "@libp2p/interface-compliance-tests/peer-discovery";
 import type { Multiaddr } from "@multiformats/multiaddr";
 import type { LightNode, PeerInfo } from "@waku/interfaces";
@@ -55,7 +56,7 @@ describe("Peer Exchange", () => {
       waku = await createLightNode();
       await waku.start();
       await waku.libp2p.dialProtocol(nwaku2Ma, PeerExchangeCodec);
-      await waitForRemotePeerWithCodec(waku, PeerExchangeCodec);
+      await waitForRemotePeerWithCodec(waku, PeerExchangeCodec, nwaku2PeerId);
 
       const components = waku.libp2p.components as unknown as Libp2pComponents;
       const peerExchange = new WakuPeerExchange(components);
@@ -77,16 +78,26 @@ describe("Peer Exchange", () => {
       expect(peerInfos[0].ENR?.peerInfo?.multiaddrs).to.not.be.null;
 
       let foundNodeMas: Multiaddr[] = [];
+      let foundNodePeerId: PeerId | undefined = undefined;
       const doesPeerIdExistInResponse =
         peerInfos.find(({ ENR }) => {
           foundNodeMas = ENR?.peerInfo?.multiaddrs ?? [];
+          foundNodePeerId = ENR?.peerInfo?.id;
           return ENR?.peerInfo?.id.toString() === nwaku1PeerId.toString();
         }) !== undefined;
 
+      if (!foundNodePeerId) {
+        throw new Error("Peer ID not found");
+      }
+
       expect(doesPeerIdExistInResponse).to.be.equal(true);
 
       await waku.libp2p.dialProtocol(foundNodeMas, PeerExchangeCodec);
-      await waitForRemotePeerWithCodec(waku, PeerExchangeCodec);
+      await waitForRemotePeerWithCodec(
+        waku,
+        PeerExchangeCodec,
+        foundNodePeerId
+      );
 
       expect(await waku.libp2p.peerStore.has(nwaku1PeerId)).to.eq(true);
       expect(waku.libp2p.getConnections()).has.length(2);