From 2e3e1e5aea35b1ee8e3e8bb2b9069c9e8dccb2a3 Mon Sep 17 00:00:00 2001
From: Alex Rudenko <OrKoN@users.noreply.github.com>
Date: Fri, 7 Jun 2024 13:39:47 +0200
Subject: [PATCH] test: add a unit test for the ws transport (#12550)

---
 .../src/node/NodeWebSocketTransport.test.ts   | 60 +++++++++++++++++++
 .../src/node/NodeWebSocketTransport.ts        | 18 +++---
 2 files changed, 68 insertions(+), 10 deletions(-)
 create mode 100644 packages/puppeteer-core/src/node/NodeWebSocketTransport.test.ts

diff --git a/packages/puppeteer-core/src/node/NodeWebSocketTransport.test.ts b/packages/puppeteer-core/src/node/NodeWebSocketTransport.test.ts
new file mode 100644
index 00000000000..e566f462aa2
--- /dev/null
+++ b/packages/puppeteer-core/src/node/NodeWebSocketTransport.test.ts
@@ -0,0 +1,60 @@
+/**
+ * @license
+ * Copyright 2024 Google Inc.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import {describe, it, beforeEach, afterEach} from 'node:test';
+
+import expect from 'expect';
+import type {WebSocket} from 'ws';
+import {WebSocketServer} from 'ws';
+
+import {NodeWebSocketTransport} from './NodeWebSocketTransport.js';
+
+describe('NodeWebSocketTransport', () => {
+  let wss: WebSocketServer;
+  let transport: NodeWebSocketTransport;
+  let connection: WebSocket;
+
+  beforeEach(async () => {
+    wss = new WebSocketServer({port: 8080});
+    wss.on('connection', c => {
+      connection = c;
+    });
+    transport = await NodeWebSocketTransport.create('ws://127.0.0.1:8080');
+  });
+
+  afterEach(() => {
+    transport.close();
+    wss.close();
+  });
+
+  it('should dispatch messages in order handling microtasks for each message first', async () => {
+    const log: string[] = [];
+    const result = new Promise<void>(resolve => {
+      transport.onmessage = message => {
+        log.push('message received ' + message);
+        return Promise.resolve().then(() => {
+          log.push('microtask1 ' + message);
+          return Promise.resolve().then(() => {
+            log.push('microtask2 ' + message);
+            if (log.length === 6) {
+              resolve();
+            }
+          });
+        });
+      };
+    });
+    connection.send('m1');
+    connection.send('m2');
+    await result;
+    expect(log).toEqual([
+      'message received m1',
+      'microtask1 m1',
+      'microtask2 m1',
+      'message received m2',
+      'microtask1 m2',
+      'microtask2 m2',
+    ]);
+  });
+});
diff --git a/packages/puppeteer-core/src/node/NodeWebSocketTransport.ts b/packages/puppeteer-core/src/node/NodeWebSocketTransport.ts
index 13f1e8349cd..d51e8130b7f 100644
--- a/packages/puppeteer-core/src/node/NodeWebSocketTransport.ts
+++ b/packages/puppeteer-core/src/node/NodeWebSocketTransport.ts
@@ -20,6 +20,8 @@ export class NodeWebSocketTransport implements ConnectionTransport {
       const ws = new NodeWebSocket(url, [], {
         followRedirects: true,
         perMessageDeflate: false,
+        // @ts-expect-error https://github.com/websockets/ws/blob/master/doc/ws.md#new-websocketaddress-protocols-options
+        allowSynchronousEvents: false,
         maxPayload: 256 * 1024 * 1024, // 256Mb
         headers: {
           'User-Agent': `Puppeteer ${packageVersion}`,
@@ -41,18 +43,14 @@ export class NodeWebSocketTransport implements ConnectionTransport {
   constructor(ws: NodeWebSocket) {
     this.#ws = ws;
     this.#ws.addEventListener('message', event => {
-      setImmediate(() => {
-        if (this.onmessage) {
-          this.onmessage.call(null, event.data);
-        }
-      });
+      if (this.onmessage) {
+        this.onmessage.call(null, event.data);
+      }
     });
     this.#ws.addEventListener('close', () => {
-      setImmediate(() => {
-        if (this.onclose) {
-          this.onclose.call(null);
-        }
-      });
+      if (this.onclose) {
+        this.onclose.call(null);
+      }
     });
     // Silently ignore all errors - we don't know what to do with them.
     this.#ws.addEventListener('error', () => {});