From 226d7d07e286ac4b67f554a68cb0e8a8aa1f5a2f Mon Sep 17 00:00:00 2001
From: Chris McCormick <chris@mccormick.cx>
Date: Mon, 29 Sep 2025 21:41:40 +0800
Subject: [PATCH] Improvements to enablePing() & tests (#506)

https://github.com/nbd-wtf/nostr-tools/pull/506
---
 pool.test.ts    | 69 +++++++++++++++++++++++++++++++++++++++----------
 relay.test.ts   | 31 ++++++++++++++++++++++
 relay.ts        |  8 +++---
 test-helpers.ts |  2 ++
 4 files changed, 93 insertions(+), 17 deletions(-)

diff --git a/pool.test.ts b/pool.test.ts
index 57ebe83..07c9a79 100644
--- a/pool.test.ts
+++ b/pool.test.ts
@@ -59,16 +59,24 @@ test('same with double subs', async () => {
   let priv = generateSecretKey()
   let pub = getPublicKey(priv)
 
-  pool.subscribeMany(relayURLs, { authors: [pub] }, {
-    onevent(event) {
-      received.push(event)
+  pool.subscribeMany(
+    relayURLs,
+    { authors: [pub] },
+    {
+      onevent(event) {
+        received.push(event)
+      },
     },
-  })
-  pool.subscribeMany(relayURLs, { authors: [pub] }, {
-    onevent(event) {
-      received.push(event)
+  )
+  pool.subscribeMany(
+    relayURLs,
+    { authors: [pub] },
+    {
+      onevent(event) {
+        received.push(event)
+      },
     },
-  })
+  )
 
   let received: Event[] = []
 
@@ -172,12 +180,16 @@ test('query a bunch of events and cancel on eose', async () => {
   let events = new Set<string>()
 
   await new Promise<void>(resolve => {
-    pool.subscribeManyEose(relayURLs, { kinds: [0, 1, 2, 3, 4, 5, 6], limit: 40 }, {
-      onevent(event) {
-        events.add(event.id)
+    pool.subscribeManyEose(
+      relayURLs,
+      { kinds: [0, 1, 2, 3, 4, 5, 6], limit: 40 },
+      {
+        onevent(event) {
+          events.add(event.id)
+        },
+        onclose: resolve as any,
       },
-      onclose: resolve as any,
-    })
+    )
   })
 
   expect(events.size).toBeGreaterThan(50)
@@ -210,6 +222,37 @@ test('get()', async () => {
   expect(event).toHaveProperty('id', ids[0])
 })
 
+test('ping-pong timeout in pool', async () => {
+  const mockRelay = mockRelays[0]
+  pool = new SimplePool({ enablePing: true })
+  const relay = await pool.ensureRelay(mockRelay.url)
+  relay.pingTimeout = 50
+  relay.pingFrequency = 50
+
+  let closed = false
+  const closedPromise = new Promise<void>(resolve => {
+    relay.onclose = () => {
+      closed = true
+      resolve()
+    }
+  })
+
+  expect(relay.connected).toBeTrue()
+
+  // wait for the first ping to succeed
+  await new Promise(resolve => setTimeout(resolve, 75))
+  expect(closed).toBeFalse()
+
+  // now make it unresponsive
+  mockRelay.unresponsive = true
+
+  // wait for the second ping to fail
+  await closedPromise
+
+  expect(relay.connected).toBeFalse()
+  expect(closed).toBeTrue()
+})
+
 test('track relays when publishing', async () => {
   let event1 = finalizeEvent(
     {
diff --git a/relay.test.ts b/relay.test.ts
index 0522297..844b4c5 100644
--- a/relay.test.ts
+++ b/relay.test.ts
@@ -117,3 +117,34 @@ test('publish timeout', async () => {
     ),
   ).rejects.toThrow('publish timed out')
 })
+
+test('ping-pong timeout', async () => {
+  const mockRelay = new MockRelay()
+  const relay = new Relay(mockRelay.url, { enablePing: true })
+  relay.pingTimeout = 50
+  relay.pingFrequency = 50
+
+  let closed = false
+  const closedPromise = new Promise<void>(resolve => {
+    relay.onclose = () => {
+      closed = true
+      resolve()
+    }
+  })
+
+  await relay.connect()
+  expect(relay.connected).toBeTrue()
+
+  // wait for the first ping to succeed
+  await new Promise(resolve => setTimeout(resolve, 75))
+  expect(closed).toBeFalse()
+
+  // now make it unresponsive
+  mockRelay.unresponsive = true
+
+  // wait for the second ping to fail
+  await closedPromise
+
+  expect(relay.connected).toBeFalse()
+  expect(closed).toBeTrue()
+})
diff --git a/relay.ts b/relay.ts
index 7539022..be6bcdb 100644
--- a/relay.ts
+++ b/relay.ts
@@ -14,12 +14,12 @@ export function useWebSocketImplementation(websocketImplementation: any) {
 }
 
 export class Relay extends AbstractRelay {
-  constructor(url: string) {
-    super(url, { verifyEvent, websocketImplementation: _WebSocket })
+  constructor(url: string, options?: { enablePing?: boolean }) {
+    super(url, { verifyEvent, websocketImplementation: _WebSocket, ...options })
   }
 
-  static async connect(url: string): Promise<Relay> {
-    const relay = new Relay(url)
+  static async connect(url: string, options?: { enablePing?: boolean }): Promise<Relay> {
+    const relay = new Relay(url, options)
     await relay.connect()
     return relay
   }
diff --git a/test-helpers.ts b/test-helpers.ts
index 9c7267e..73c509e 100644
--- a/test-helpers.ts
+++ b/test-helpers.ts
@@ -26,6 +26,7 @@ export class MockRelay {
   public url: string
   public secretKeys: Uint8Array[]
   public preloadedEvents: Event[]
+  public unresponsive: boolean = false
 
   constructor(url?: string | undefined) {
     serial++
@@ -48,6 +49,7 @@ export class MockRelay {
       let subs: { [subId: string]: { conn: any; filters: Filter[] } } = {}
 
       conn.on('message', (message: string) => {
+        if (this.unresponsive) return
         const data = JSON.parse(message)
 
         switch (data[0]) {