diff --git a/packages/playwright/src/isomorphic/testServerConnection.ts b/packages/playwright/src/isomorphic/testServerConnection.ts
index 34ce0e31ef..2a77a8307e 100644
--- a/packages/playwright/src/isomorphic/testServerConnection.ts
+++ b/packages/playwright/src/isomorphic/testServerConnection.ts
@@ -19,7 +19,9 @@ import * as events from './events';
 
 // -- Reuse boundary -- Everything below this line is reused in the vscode extension.
 
-export class TestServerConnection implements TestServerInterface, TestServerInterfaceEvents {
+export type TestServerSocket = Pick<WebSocket, 'addEventListener' | 'send' | 'close'>;
+
+abstract class TestServerConnection implements TestServerInterface, TestServerInterfaceEvents {
   readonly onClose: events.Event<void>;
   readonly onReport: events.Event<any>;
   readonly onStdio: events.Event<{ type: 'stderr' | 'stdout'; text?: string | undefined; buffer?: string | undefined; }>;
@@ -33,20 +35,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 +66,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 +87,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 +202,14 @@ export class TestServerConnection implements TestServerInterface, TestServerInte
 
   close() {
     try {
-      this._ws.close();
+      this._socket.close();
     } catch {
     }
   }
 }
+
+export class WSTestServerConnection extends TestServerConnection {
+  constructor(wsURL: string) {
+    super(new WebSocket(wsURL));
+  }
+}
diff --git a/packages/trace-viewer/src/ui/uiModeTestListView.tsx b/packages/trace-viewer/src/ui/uiModeTestListView.tsx
index b7a9d8e1d1..77edba2249 100644
--- a/packages/trace-viewer/src/ui/uiModeTestListView.tsx
+++ b/packages/trace-viewer/src/ui/uiModeTestListView.tsx
@@ -29,7 +29,7 @@ import type { SourceLocation } from './modelUtil';
 import { testStatusIcon } from './testUtils';
 import type { TestModel } from './uiModeModel';
 import './uiModeTestListView.css';
-import type { TestServerConnection } from '@testIsomorphic/testServerConnection';
+import type { WSTestServerConnection } from '@testIsomorphic/testServerConnection';
 import { TagView } from './tag';
 
 const TestTreeView = TreeView<TreeItem>;
@@ -37,7 +37,7 @@ const TestTreeView = TreeView<TreeItem>;
 export const TestListView: React.FC<{
   filterText: string,
   testTree: TestTree,
-  testServerConnection: TestServerConnection | undefined,
+  testServerConnection: WSTestServerConnection | undefined,
   testModel?: TestModel,
   runTests: (mode: 'bounce-if-busy' | 'queue-if-busy', testIds: Set<string>) => void,
   runningState?: { testIds: Set<string>, itemSelectedByUser?: boolean, completed?: boolean },
diff --git a/packages/trace-viewer/src/ui/uiModeView.tsx b/packages/trace-viewer/src/ui/uiModeView.tsx
index b12617c300..7b3e13cf18 100644
--- a/packages/trace-viewer/src/ui/uiModeView.tsx
+++ b/packages/trace-viewer/src/ui/uiModeView.tsx
@@ -33,7 +33,7 @@ import { useDarkModeSetting } from '@web/theme';
 import { clsx, settings, useSetting } from '@web/uiUtils';
 import { statusEx, TestTree } from '@testIsomorphic/testTree';
 import type { TreeItem  } from '@testIsomorphic/testTree';
-import { TestServerConnection } from '@testIsomorphic/testServerConnection';
+import { WSTestServerConnection } from '@testIsomorphic/testServerConnection';
 import { pathSeparator } from './uiModeModel';
 import type { TestModel } from './uiModeModel';
 import { FiltersView } from './uiModeFiltersView';
@@ -95,7 +95,7 @@ export const UIModeView: React.FC<{}> = ({
   const [collapseAllCount, setCollapseAllCount] = React.useState(0);
   const [isDisconnected, setIsDisconnected] = React.useState(false);
   const [hasBrowsers, setHasBrowsers] = React.useState(true);
-  const [testServerConnection, setTestServerConnection] = React.useState<TestServerConnection>();
+  const [testServerConnection, setTestServerConnection] = React.useState<WSTestServerConnection>();
   const [teleSuiteUpdater, setTeleSuiteUpdater] = React.useState<TeleSuiteUpdater>();
   const [settingsVisible, setSettingsVisible] = React.useState(false);
   const [testingOptionsVisible, setTestingOptionsVisible] = React.useState(false);
@@ -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 WSTestServerConnection(wsURL.toString()));
   }, []);
 
   // Load tests on startup.
diff --git a/packages/trace-viewer/src/ui/workbenchLoader.tsx b/packages/trace-viewer/src/ui/workbenchLoader.tsx
index e7f94e969a..cac10727bc 100644
--- a/packages/trace-viewer/src/ui/workbenchLoader.tsx
+++ b/packages/trace-viewer/src/ui/workbenchLoader.tsx
@@ -21,7 +21,7 @@ import { MultiTraceModel } from './modelUtil';
 import './workbenchLoader.css';
 import { toggleTheme } from '@web/theme';
 import { Workbench } from './workbench';
-import { TestServerConnection } from '@testIsomorphic/testServerConnection';
+import { WSTestServerConnection } from '@testIsomorphic/testServerConnection';
 
 export const WorkbenchLoader: React.FunctionComponent<{
 }> = () => {
@@ -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 WSTestServerConnection(wsURL.toString());
       testServerConnection.onLoadTraceRequested(async params => {
         setTraceURLs(params.traceUrl ? [params.traceUrl] : []);
         setDragOver(false);
diff --git a/tests/playwright-test/test-server-connection.spec.ts b/tests/playwright-test/test-server-connection.spec.ts
index af5f0223e5..0c1afd00ad 100644
--- a/tests/playwright-test/test-server-connection.spec.ts
+++ b/tests/playwright-test/test-server-connection.spec.ts
@@ -15,9 +15,9 @@
  */
 import { test as baseTest, expect } from './ui-mode-fixtures';
 
-import { TestServerConnection } from '../../packages/playwright/lib/isomorphic/testServerConnection';
+import { WSTestServerConnection } from '../../packages/playwright/lib/isomorphic/testServerConnection';
 
-class TestServerConnectionUnderTest extends TestServerConnection {
+class TestServerConnectionUnderTest extends WSTestServerConnection {
   events: [string, any][] = [];
 
   constructor(wsUrl: string) {