From 8ca1f63b0769464c6955bea2a64522ff22841370 Mon Sep 17 00:00:00 2001
From: Simon Knott <info@simonknott.de>
Date: Thu, 22 Aug 2024 16:56:10 +0200
Subject: [PATCH] chore(ui): decouple TestServerConnection from websocket
 transport

---
 .../src/isomorphic/testServerConnection.ts    | 30 ++++++++++++-------
 packages/trace-viewer/src/ui/uiModeView.tsx   |  2 +-
 .../trace-viewer/src/ui/workbenchLoader.tsx   |  2 +-
 3 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/packages/playwright/src/isomorphic/testServerConnection.ts b/packages/playwright/src/isomorphic/testServerConnection.ts
index 34ce0e31ef..9c5596e4ee 100644
--- a/packages/playwright/src/isomorphic/testServerConnection.ts
+++ b/packages/playwright/src/isomorphic/testServerConnection.ts
@@ -19,6 +19,16 @@ import * as events from './events';
 
 // -- Reuse boundary -- Everything below this line is reused in the vscode extension.
 
+export interface TestServerSocket {
+  addEventListener(type: 'message', listener: (event: { data: string }) => void): void;
+  addEventListener(type: 'open', listener: () => void): void;
+  addEventListener(type: 'error', listener: () => void): void;
+  addEventListener(type: 'close', listener: () => void): void;
+
+  send(data: string): void;
+  close(): void;
+}
+
 export class TestServerConnection implements TestServerInterface, TestServerInterfaceEvents {
   readonly onClose: events.Event<void>;
   readonly onReport: events.Event<any>;
@@ -33,20 +43,20 @@ export class TestServerConnection implements TestServerInterface, TestServerInte
   private _onLoadTraceRequestedEmitter = new events.EventEmitter<{ traceUrl: string }>();
 
   private _lastId = 0;
-  private _ws: WebSocket;
+  private _socket: TestServerSocket;
   private _callbacks = new Map<number, { resolve: (arg: any) => void, reject: (arg: Error) => void }>();
   private _connectedPromise: Promise<void>;
   private _isClosed = false;
 
-  constructor(wsURL: string) {
+  constructor(socket: TestServerSocket) {
     this.onClose = this._onCloseEmitter.event;
     this.onReport = this._onReportEmitter.event;
     this.onStdio = this._onStdioEmitter.event;
     this.onTestFilesChanged = this._onTestFilesChangedEmitter.event;
     this.onLoadTraceRequested = this._onLoadTraceRequestedEmitter.event;
 
-    this._ws = new WebSocket(wsURL);
-    this._ws.addEventListener('message', event => {
+    this._socket = socket;
+    this._socket.addEventListener('message', event => {
       const message = JSON.parse(String(event.data));
       const { id, result, error, method, params } = message;
       if (id) {
@@ -64,10 +74,10 @@ export class TestServerConnection implements TestServerInterface, TestServerInte
     });
     const pingInterval = setInterval(() => this._sendMessage('ping').catch(() => {}), 30000);
     this._connectedPromise = new Promise<void>((f, r) => {
-      this._ws.addEventListener('open', () => f());
-      this._ws.addEventListener('error', r);
+      this._socket.addEventListener('open', () => f());
+      this._socket.addEventListener('error', r);
     });
-    this._ws.addEventListener('close', () => {
+    this._socket.addEventListener('close', () => {
       this._isClosed = true;
       this._onCloseEmitter.fire();
       clearInterval(pingInterval);
@@ -85,7 +95,7 @@ export class TestServerConnection implements TestServerInterface, TestServerInte
     await this._connectedPromise;
     const id = ++this._lastId;
     const message = { id, method, params };
-    this._ws.send(JSON.stringify(message));
+    this._socket.send(JSON.stringify(message));
     return new Promise((resolve, reject) => {
       this._callbacks.set(id, { resolve, reject });
     });
@@ -200,8 +210,8 @@ export class TestServerConnection implements TestServerInterface, TestServerInte
 
   close() {
     try {
-      this._ws.close();
+      this._socket.close();
     } catch {
     }
   }
-}
+}
\ No newline at end of file
diff --git a/packages/trace-viewer/src/ui/uiModeView.tsx b/packages/trace-viewer/src/ui/uiModeView.tsx
index b12617c300..2525ec4b35 100644
--- a/packages/trace-viewer/src/ui/uiModeView.tsx
+++ b/packages/trace-viewer/src/ui/uiModeView.tsx
@@ -135,7 +135,7 @@ export const UIModeView: React.FC<{}> = ({
   const inputRef = React.useRef<HTMLInputElement>(null);
 
   const reloadTests = React.useCallback(() => {
-    setTestServerConnection(new TestServerConnection(wsURL.toString()));
+    setTestServerConnection(new TestServerConnection(new WebSocket(wsURL)));
   }, []);
 
   // Load tests on startup.
diff --git a/packages/trace-viewer/src/ui/workbenchLoader.tsx b/packages/trace-viewer/src/ui/workbenchLoader.tsx
index e7f94e969a..3291008066 100644
--- a/packages/trace-viewer/src/ui/workbenchLoader.tsx
+++ b/packages/trace-viewer/src/ui/workbenchLoader.tsx
@@ -102,7 +102,7 @@ export const WorkbenchLoader: React.FunctionComponent<{
       const guid = new URLSearchParams(window.location.search).get('ws');
       const wsURL = new URL(`../${guid}`, window.location.toString());
       wsURL.protocol = (window.location.protocol === 'https:' ? 'wss:' : 'ws:');
-      const testServerConnection = new TestServerConnection(wsURL.toString());
+      const testServerConnection = new TestServerConnection(new WebSocket(wsURL));
       testServerConnection.onLoadTraceRequested(async params => {
         setTraceURLs(params.traceUrl ? [params.traceUrl] : []);
         setDragOver(false);