From 4f58bde054d2dddaa8cc2df1f4d97c4ef293d416 Mon Sep 17 00:00:00 2001
From: Franck Royer <franck@royer.one>
Date: Tue, 23 Mar 2021 12:14:48 +1100
Subject: [PATCH] Add js to js subscribe test

---
 src/lib/waku.ts            |  2 +-
 src/lib/waku_relay.spec.ts | 80 +++++++++++++++++++-------------------
 2 files changed, 42 insertions(+), 40 deletions(-)

diff --git a/src/lib/waku.ts b/src/lib/waku.ts
index ef1f7738c5..f9c3e0c9f7 100644
--- a/src/lib/waku.ts
+++ b/src/lib/waku.ts
@@ -39,7 +39,7 @@ export default class Waku {
 
   async dialWithMultiAddr(peerId: PeerId, multiaddr: Multiaddr[]) {
     this.libp2p.peerStore.addressBook.set(peerId, multiaddr);
-    return this.libp2p.dialProtocol(peerId, CODEC);
+    await this.libp2p.dialProtocol(peerId, CODEC);
   }
 
   async stop() {
diff --git a/src/lib/waku_relay.spec.ts b/src/lib/waku_relay.spec.ts
index f28dc3ba87..b63dc49444 100644
--- a/src/lib/waku_relay.spec.ts
+++ b/src/lib/waku_relay.spec.ts
@@ -10,58 +10,60 @@ import { Message } from './waku_message';
 import { CODEC, TOPIC } from './waku_relay';
 
 describe('Waku Relay', () => {
-  // TODO: Fix this, see https://github.com/ChainSafe/js-libp2p-gossipsub/issues/151
-  it.skip('Publish', async () => {
-    const message = Message.fromUtf8String('Bird bird bird, bird is the word!');
-
-    const [waku1, waku2] = await Promise.all([
+  let waku1: Waku;
+  let waku2: Waku;
+  beforeEach(async function () {
+    [waku1, waku2] = await Promise.all([
       Waku.create(NOISE_KEY_1),
       Waku.create(NOISE_KEY_2),
     ]);
 
-    // Add node's 2 data to the PeerStore
-    waku1.libp2p.peerStore.addressBook.set(
-      waku2.libp2p.peerId,
-      waku2.libp2p.multiaddrs
-    );
-    await waku1.libp2p.dial(waku2.libp2p.peerId);
+    await waku1.dialWithMultiAddr(waku2.libp2p.peerId, waku2.libp2p.multiaddrs);
 
+    await Promise.all([
+      new Promise((resolve) =>
+        waku1.libp2p.pubsub.once('gossipsub:heartbeat', resolve)
+      ),
+      new Promise((resolve) =>
+        waku2.libp2p.pubsub.once('gossipsub:heartbeat', resolve)
+      ),
+    ]);
+
+    await waku1.relay.subscribe();
     await waku2.relay.subscribe();
-    await new Promise((resolve) =>
-      waku2.libp2p.pubsub.once('pubsub:subscription-change', (...args) =>
-        resolve(args)
-      )
-    );
 
-    // Setup the promise before publishing to ensure the event is not missed
-    const promise = waitForNextData(waku1.libp2p.pubsub);
-
-    await waku2.relay.publish(message);
-
-    const node1Received = await promise;
-
-    expect(node1Received.isEqualTo(message)).to.be.true;
-
-    await Promise.all([waku1.stop(), waku2.stop()]);
+    await Promise.all([
+      new Promise((resolve) =>
+        waku1.libp2p.pubsub.once('pubsub:subscription-change', (...args) =>
+          resolve(args)
+        )
+      ),
+      new Promise((resolve) =>
+        waku2.libp2p.pubsub.once('pubsub:subscription-change', (...args) =>
+          resolve(args)
+        )
+      ),
+    ]);
   });
 
-  it('Registers waku relay protocol', async function () {
-    const waku = await Waku.create(NOISE_KEY_1);
+  afterEach(async function () {
+    await waku1.stop();
+    await waku2.stop();
+  });
 
-    const protocols = Array.from(waku.libp2p.upgrader.protocols.keys());
+  it('Subscribe', async function () {
+    const subscribers1 = waku1.libp2p.pubsub.getSubscribers(TOPIC);
+    const subscribers2 = waku2.libp2p.pubsub.getSubscribers(TOPIC);
+
+    expect(subscribers1).to.contain(waku2.libp2p.peerId.toB58String());
+    expect(subscribers2).to.contain(waku1.libp2p.peerId.toB58String());
+  });
+
+  it('Register correct protocols', async function () {
+    const protocols = Array.from(waku1.libp2p.upgrader.protocols.keys());
 
     expect(protocols).to.contain(CODEC);
-
-    await waku.stop();
-  });
-
-  it('Does not register any sub protocol', async function () {
-    const waku = await Waku.create(NOISE_KEY_1);
-
-    const protocols = Array.from(waku.libp2p.upgrader.protocols.keys());
     expect(protocols.findIndex((value) => value.match(/sub/))).to.eq(-1);
-
-    await waku.stop();
   });
 
   describe('Interop: Nim', function () {