test(websocket): fully deflake send test

test/websocket/send.js

@@ -1,131 +1,152 @@
 'use strict'
 
+const { once } = require('node:events')
 const { test, describe } = require('node:test')
 const { WebSocketServer } = require('ws')
 const { WebSocket } = require('../..')
 
+async function closeWebSocket (ws) {
+  if (ws.readyState === WebSocket.CLOSED) {
+    return
+  }
+
+  const close = once(ws, 'close')
+  ws.close()
+  await close
+}
+
+async function terminateWebSocket (ws) {
+  if (ws.readyState === WebSocket.CLOSED) {
+    return
+  }
+
+  const close = once(ws, 'close')
+  ws.terminate()
+  await close
+}
+
+async function closeWebSocketServer (server) {
+  await Promise.allSettled(Array.from(server.clients, terminateWebSocket))
+  await new Promise((resolve) => server.close(resolve))
+}
+
+function registerCleanup (t, ws, server) {
+  t.after(async () => {
+    await Promise.allSettled([
+      closeWebSocket(ws),
+      closeWebSocketServer(server)
+    ])
+  })
+}
+
 // the following three tests exercise different code paths because of the three
 // different ways a payload length may be specified in a WebSocket frame
 // (https://datatracker.ietf.org/doc/html/rfc6455#section-5.2)
 
-test('Sending >= 2^16 bytes', (t) => {
+test('Sending >= 2^16 bytes', async (t) => {
   const server = new WebSocketServer({ port: 0 })
 
-  server.on('connection', (ws) => {
-    ws.on('message', (m, isBinary) => {
-      ws.send(m, { binary: isBinary })
+  server.on('connection', (socket) => {
+    socket.on('message', (message, isBinary) => {
+      socket.send(message, { binary: isBinary })
     })
   })
 
   const payload = Buffer.allocUnsafe(2 ** 16).fill('Hello')
 
   const ws = new WebSocket(`ws://localhost:${server.address().port}`)
+  registerCleanup(t, ws, server)
 
-  ws.addEventListener('open', () => {
-    ws.send(payload)
-  })
-
-  return new Promise((resolve) => {
-    ws.addEventListener('message', async ({ data }) => {
-      t.assert.ok(data instanceof Blob)
-      t.assert.strictEqual(data.size, payload.length)
-      t.assert.deepStrictEqual(Buffer.from(await data.arrayBuffer()), payload)
+  await once(ws, 'open')
+  ws.send(payload)
 
-      ws.close()
-      server.close()
+  const [{ data }] = await once(ws, 'message')
 
-      resolve()
-    })
-  })
+  t.assert.ok(data instanceof Blob)
+  t.assert.strictEqual(data.size, payload.length)
+  t.assert.deepStrictEqual(Buffer.from(await data.arrayBuffer()), payload)
 })
 
-test('Sending >= 126, < 2^16 bytes', (t) => {
+test('Sending >= 126, < 2^16 bytes', async (t) => {
   const server = new WebSocketServer({ port: 0 })
 
-  server.on('connection', (ws) => {
-    ws.on('message', (m, isBinary) => {
-      ws.send(m, { binary: isBinary })
+  server.on('connection', (socket) => {
+    socket.on('message', (message, isBinary) => {
+      socket.send(message, { binary: isBinary })
     })
   })
 
   const payload = Buffer.allocUnsafe(126).fill('Hello')
 
   const ws = new WebSocket(`ws://localhost:${server.address().port}`)
+  registerCleanup(t, ws, server)
 
-  ws.addEventListener('open', () => {
-    ws.send(payload)
-  })
+  await once(ws, 'open')
+  ws.send(payload)
 
-  return new Promise((resolve) => {
-    ws.addEventListener('message', async ({ data }) => {
-      t.assert.ok(data instanceof Blob)
-      t.assert.strictEqual(data.size, payload.length)
-      t.assert.deepStrictEqual(Buffer.from(await data.arrayBuffer()), payload)
+  const [{ data }] = await once(ws, 'message')
 
-      ws.close()
-      server.close()
-      resolve()
-    })
-  })
+  t.assert.ok(data instanceof Blob)
+  t.assert.strictEqual(data.size, payload.length)
+  t.assert.deepStrictEqual(Buffer.from(await data.arrayBuffer()), payload)
 })
 
-test('Sending < 126 bytes', (t) => {
+test('Sending < 126 bytes', async (t) => {
   const server = new WebSocketServer({ port: 0 })
 
-  server.on('connection', (ws) => {
-    ws.on('message', (m, isBinary) => {
-      ws.send(m, { binary: isBinary })
+  server.on('connection', (socket) => {
+    socket.on('message', (message, isBinary) => {
+      socket.send(message, { binary: isBinary })
     })
   })
 
   const payload = Buffer.allocUnsafe(125).fill('Hello')
 
   const ws = new WebSocket(`ws://localhost:${server.address().port}`)
+  registerCleanup(t, ws, server)
 
-  ws.addEventListener('open', () => {
-    ws.send(payload)
-  })
+  await once(ws, 'open')
+  ws.send(payload)
 
-  return new Promise((resolve) => {
-    ws.addEventListener('message', async ({ data }) => {
-      t.assert.ok(data instanceof Blob)
-      t.assert.strictEqual(data.size, payload.length)
-      t.assert.deepStrictEqual(Buffer.from(await data.arrayBuffer()), payload)
+  const [{ data }] = await once(ws, 'message')
 
-      ws.close()
-      server.close()
-      resolve()
-    })
-  })
+  t.assert.ok(data instanceof Blob)
+  t.assert.strictEqual(data.size, payload.length)
+  t.assert.deepStrictEqual(Buffer.from(await data.arrayBuffer()), payload)
 })
 
-test('Sending data after close', (t) => {
+test('Sending data after close', async (t) => {
   const server = new WebSocketServer({ port: 0 })
-
   const ws = new WebSocket(`ws://localhost:${server.address().port}`)
+  registerCleanup(t, ws, server)
 
-  return new Promise((resolve, reject) => {
-    server.on('connection', (ws) => {
-      ws.on('message', reject)
-    })
+  const connection = once(server, 'connection')
 
-    ws.addEventListener('open', () => {
-      ws.close()
-      ws.send('Some message')
-      server.close()
-
-      resolve()
-    })
+  await once(ws, 'open')
 
-    ws.addEventListener('error', reject)
+  const [socket] = await connection
+  socket.on('message', () => {
+    t.assert.fail('Received unexpected message after closing the client')
   })
+
+  const clientClose = once(ws, 'close')
+  const socketClose = once(socket, 'close')
+
+  ws.close()
+  ws.send('Some message')
+
+  await Promise.all([
+    clientClose,
+    socketClose
+  ])
 })
 
 test('Sending data before connected', (t) => {
   const server = new WebSocketServer({ port: 0 })
-
   const ws = new WebSocket(`ws://localhost:${server.address().port}`)
 
+  registerCleanup(t, ws, server)
+
   t.assert.throws(
     () => ws.send('Not sent'),
     {
@@ -135,101 +156,80 @@ test('Sending data before connected', (t) => {
   )
 
   t.assert.strictEqual(ws.readyState, WebSocket.CONNECTING)
-  server.close()
 })
 
 describe('Sending data to a server', () => {
-  test('Send with string', (t) => {
+  test('Send with string', async (t) => {
     const server = new WebSocketServer({ port: 0 })
-
     const ws = new WebSocket(`ws://localhost:${server.address().port}`)
+    registerCleanup(t, ws, server)
 
-    ws.addEventListener('open', () => {
-      ws.send('message')
-    })
+    const connection = once(server, 'connection')
 
-    return new Promise((resolve) => {
-      server.on('connection', (ws) => {
-        ws.on('message', (data, isBinary) => {
-          t.assert.ok(!isBinary, 'Received text frame')
-          t.assert.deepStrictEqual(data, Buffer.from('message'))
-          ws.close(1000)
-          server.close()
-          resolve()
-        })
-      })
-    })
+    await once(ws, 'open')
+    ws.send('message')
+
+    const [socket] = await connection
+    const [data, isBinary] = await once(socket, 'message')
+
+    t.assert.ok(!isBinary, 'Received text frame')
+    t.assert.deepStrictEqual(data, Buffer.from('message'))
   })
 
-  test('Send with ArrayBuffer', (t) => {
+  test('Send with ArrayBuffer', async (t) => {
     const message = new TextEncoder().encode('message')
     const ab = new ArrayBuffer(7)
     new Uint8Array(ab).set(message)
 
     const server = new WebSocketServer({ port: 0 })
-
     const ws = new WebSocket(`ws://localhost:${server.address().port}`)
+    registerCleanup(t, ws, server)
 
-    ws.addEventListener('open', () => {
-      ws.send(ab)
-    })
+    const connection = once(server, 'connection')
 
-    return new Promise((resolve) => {
-      server.on('connection', (ws) => {
-        ws.on('message', (data, isBinary) => {
-          t.assert.ok(isBinary)
-          t.assert.deepStrictEqual(new Uint8Array(data), message)
-          ws.close(1000)
-          server.close()
-          resolve()
-        })
-      })
-    })
+    await once(ws, 'open')
+    ws.send(ab)
+
+    const [socket] = await connection
+    const [data, isBinary] = await once(socket, 'message')
+
+    t.assert.ok(isBinary)
+    t.assert.deepStrictEqual(new Uint8Array(data), message)
   })
 
-  test('Send with Blob', (t) => {
+  test('Send with Blob', async (t) => {
     const blob = new Blob(['hello'])
     const server = new WebSocketServer({ port: 0 })
-
     const ws = new WebSocket(`ws://localhost:${server.address().port}`)
+    registerCleanup(t, ws, server)
 
-    ws.addEventListener('open', () => {
-      ws.send(blob)
-    })
+    const connection = once(server, 'connection')
 
-    return new Promise((resolve) => {
-      server.on('connection', (ws) => {
-        ws.on('message', (data, isBinary) => {
-          t.assert.ok(isBinary)
-          t.assert.deepStrictEqual(data, Buffer.from('hello'))
-          ws.close(1000)
-          server.close()
-          resolve()
-        })
-      })
-    })
+    await once(ws, 'open')
+    ws.send(blob)
+
+    const [socket] = await connection
+    const [data, isBinary] = await once(socket, 'message')
+
+    t.assert.ok(isBinary)
+    t.assert.deepStrictEqual(data, Buffer.from('hello'))
   })
 
-  test('Cannot send with SharedArrayBuffer', (t) => {
+  test('Cannot send with SharedArrayBuffer', async (t) => {
     const sab = new SharedArrayBuffer(0)
     const server = new WebSocketServer({ port: 0 })
-
     const ws = new WebSocket(`ws://localhost:${server.address().port}`)
+    registerCleanup(t, ws, server)
 
-    ws.addEventListener('open', () => {
-      ws.send(sab)
-    })
+    const connection = once(server, 'connection')
 
-    return new Promise((resolve) => {
-      server.on('connection', (ws) => {
-        ws.on('message', (data, isBinary) => {
-          t.assert.ok(!isBinary)
-          t.assert.deepStrictEqual(data, Buffer.from('[object SharedArrayBuffer]'))
-          ws.close(1000)
-          server.close()
-          resolve()
-        })
-      })
-    })
+    await once(ws, 'open')
+    ws.send(sab)
+
+    const [socket] = await connection
+    const [data, isBinary] = await once(socket, 'message')
+
+    t.assert.ok(!isBinary)
+    t.assert.deepStrictEqual(data, Buffer.from('[object SharedArrayBuffer]'))
   })
 })