test: deflake invalid websocket close status test

test/websocket/close-invalid-status-code.js

@@ -5,32 +5,33 @@ const { once } = require('node:events')
 const { WebSocketServer } = require('ws')
 const { WebSocket } = require('../..')
 
-test('Client fails the connection if receiving a masked frame', async (t) => {
-  t.plan(2)
-
+test('Client fails the connection if receiving an invalid close status code', async (t) => {
+  const invalidCloseFrame = Buffer.from([0x88, 0x02, 0x03, 0xee])
   const server = new WebSocketServer({ port: 0 })
 
-  server.on('connection', (ws) => {
-    const socket = ws._socket
+  t.after(() => {
+    for (const client of server.clients) {
+      client.terminate()
+    }
 
-    // 1006 status code
-    socket.write(Buffer.from([0x88, 0x02, 0x03, 0xee]), () => ws.close())
+    server.close()
   })
 
-  const ws = new WebSocket(`ws://localhost:${server.address().port}`)
+  await once(server, 'listening')
 
-  ws.addEventListener('close', (e) => {
-    t.assert.deepStrictEqual(e.code, 1006)
+  server.on('connection', (serverWs) => {
+    // 1006 status code
+    serverWs._socket.end(invalidCloseFrame)
   })
 
-  ws.addEventListener('error', () => {
-    t.assert.ok(true)
-  })
+  const ws = new WebSocket(`ws://127.0.0.1:${server.address().port}`)
+  t.after(() => ws.close())
 
-  t.after(() => {
-    server.close()
-    ws.close()
-  })
+  const [[errorEvent], [closeEvent]] = await Promise.all([
+    once(ws, 'error'),
+    once(ws, 'close')
+  ])
 
-  await once(ws, 'close')
+  t.assert.ok(errorEvent)
+  t.assert.strictEqual(closeEvent.code, 1006)
 })