From dc96074c810c145ea7eaafc215b498206b3b11f3 Mon Sep 17 00:00:00 2001
From: Danish Arora <35004822+danisharora099@users.noreply.github.com>
Date: Thu, 11 Jan 2024 16:39:48 +0530
Subject: [PATCH] fix(tests): peer-exchange interop (#1773)

* fix(tests): px interop

* optimize waitFor utility

* address minor comments

* fix: test refactoring
---
 packages/tests/src/index.ts                   |  1 +
 .../tests/src/waitForRemotePeerWithCodec.ts   | 39 +++++++++++++++++++
 .../tests/tests/peer_exchange.node.spec.ts    | 38 ++++++++++++++----
 3 files changed, 71 insertions(+), 7 deletions(-)
 create mode 100644 packages/tests/src/waitForRemotePeerWithCodec.ts

diff --git a/packages/tests/src/index.ts b/packages/tests/src/index.ts
index d56386239b..43901b03af 100644
--- a/packages/tests/src/index.ts
+++ b/packages/tests/src/index.ts
@@ -13,3 +13,4 @@ export * from "./node/node.js";
 export * from "./teardown.js";
 export * from "./message_collector.js";
 export * from "./utils.js";
+export * from "./waitForRemotePeerWithCodec.js";
diff --git a/packages/tests/src/waitForRemotePeerWithCodec.ts b/packages/tests/src/waitForRemotePeerWithCodec.ts
new file mode 100644
index 0000000000..411b6116dd
--- /dev/null
+++ b/packages/tests/src/waitForRemotePeerWithCodec.ts
@@ -0,0 +1,39 @@
+import type { IdentifyResult } from "@libp2p/interface";
+import type { PeerId } from "@libp2p/interface/peer-id";
+import type { LightNode } from "@waku/interfaces";
+
+/**
+ * Wait for a remote peer to be identified with a given codec
+ * @param waku - Waku node
+ * @param codec - Codec to wait for
+ * @returns Promise that resolves when the peer is identified
+ * @internal
+ * This function is introduced as `core/waitForRemotePeer` only accounts for core protocols like Filter, LightPush & Store
+ * While this (currently) is not required by the SDK, it is required for the tests
+ */
+export async function waitForRemotePeerWithCodec(
+  waku: LightNode,
+  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)) {
+        waku.libp2p.removeEventListener("peer:identify", cb);
+        resolve();
+      }
+    };
+    waku.libp2p.addEventListener("peer:identify", cb);
+  });
+
+  return;
+}
diff --git a/packages/tests/tests/peer_exchange.node.spec.ts b/packages/tests/tests/peer_exchange.node.spec.ts
index 657d2c1781..996e1e1adc 100644
--- a/packages/tests/tests/peer_exchange.node.spec.ts
+++ b/packages/tests/tests/peer_exchange.node.spec.ts
@@ -1,4 +1,6 @@
+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";
 import {
   PeerExchangeCodec,
@@ -9,7 +11,7 @@ import { createLightNode, Libp2pComponents } from "@waku/sdk";
 import { expect } from "chai";
 
 import { delay } from "../src/delay.js";
-import { tearDownNodes } from "../src/index.js";
+import { tearDownNodes, waitForRemotePeerWithCodec } from "../src/index.js";
 import { makeLogFileName } from "../src/log_file.js";
 import { NimGoNode } from "../src/node/node.js";
 
@@ -48,11 +50,13 @@ describe("Peer Exchange", () => {
       });
 
       const nwaku1PeerId = await nwaku1.getPeerId();
+      const nwaku2PeerId = await nwaku2.getPeerId();
       const nwaku2Ma = await nwaku2.getMultiaddrWithId();
 
       waku = await createLightNode();
       await waku.start();
       await waku.libp2p.dialProtocol(nwaku2Ma, PeerExchangeCodec);
+      await waitForRemotePeerWithCodec(waku, PeerExchangeCodec, nwaku2PeerId);
 
       const components = waku.libp2p.components as unknown as Libp2pComponents;
       const peerExchange = new WakuPeerExchange(components);
@@ -62,6 +66,7 @@ describe("Peer Exchange", () => {
       let peerInfos: PeerInfo[] = [];
       while (peerInfos.length <= 0) {
         peerInfos = (await peerExchange.query({
+          peerId: nwaku2PeerId,
           numPeers: numPeersToRequest
         })) as PeerInfo[];
         await delay(3000);
@@ -70,16 +75,35 @@ describe("Peer Exchange", () => {
       expect(peerInfos.length).to.be.greaterThan(0);
       expect(peerInfos.length).to.be.lessThanOrEqual(numPeersToRequest);
       expect(peerInfos[0].ENR).to.not.be.null;
+      expect(peerInfos[0].ENR?.peerInfo?.multiaddrs).to.not.be.null;
 
-      const doesPeerIdExistInResponse =
-        peerInfos.find(
-          ({ ENR }) => ENR?.peerInfo?.id.toString() === nwaku1PeerId.toString()
-        ) !== undefined;
+      let foundNodeMultiaddrs: Multiaddr[] = [];
+      let foundNodePeerId: PeerId | undefined = undefined;
+      const doesPeerIdExistInResponse = peerInfos.some(({ ENR }) => {
+        foundNodeMultiaddrs = ENR?.peerInfo?.multiaddrs ?? [];
+        foundNodePeerId = ENR?.peerInfo?.id;
+        return ENR?.peerInfo?.id.toString() === nwaku1PeerId.toString();
+      });
+
+      if (!foundNodePeerId) {
+        throw new Error("Peer ID not found");
+      }
+
+      if (!foundNodePeerId) {
+        throw new Error("Peer ID not found");
+      }
 
       expect(doesPeerIdExistInResponse).to.be.equal(true);
 
-      expect(await waku.libp2p.peerStore.has(await nwaku2.getPeerId())).to.be
-        .true;
+      await waku.libp2p.dialProtocol(foundNodeMultiaddrs, PeerExchangeCodec);
+      await waitForRemotePeerWithCodec(
+        waku,
+        PeerExchangeCodec,
+        foundNodePeerId
+      );
+
+      expect(await waku.libp2p.peerStore.has(nwaku1PeerId)).to.eq(true);
+      expect(waku.libp2p.getConnections()).has.length(2);
     });
   });