fix(ui-mode): prevent websocket connection leaks on reload (#33643)

packages/trace-viewer/src/ui/uiModeView.tsx

@@ -109,7 +109,10 @@ export const UIModeView: React.FC<{}> = ({
   const inputRef = React.useRef<HTMLInputElement>(null);
 
   const reloadTests = React.useCallback(() => {
-    setTestServerConnection(new TestServerConnection(new WebSocketTestServerTransport(wsURL)));
+    setTestServerConnection(prevConnection => {
+      prevConnection?.close();
+      return new TestServerConnection(new WebSocketTestServerTransport(wsURL));
+    });
   }, []);
 
   // Load tests on startup.
@@ -224,7 +227,7 @@ export const UIModeView: React.FC<{}> = ({
         newFilter.set(projectSuite.title, !!selectedProjects?.includes(projectSuite.title));
     }
     if (!selectedProjects && newFilter.size && ![...newFilter.values()].includes(true))
-      newFilter.set(newFilter.entries().next().value[0], true);
+      newFilter.set(newFilter.entries().next().value![0], true);
     if (projectFilters.size !== newFilter.size || [...projectFilters].some(([k, v]) => newFilter.get(k) !== v))
       setProjectFilters(newFilter);
   }, [projectFilters, testModel]);

tests/playwright-test/ui-mode-test-run.spec.ts

@@ -775,3 +775,26 @@ test('should respect --ignore-snapshots option', {
           - treeitem ${/\[icon-check\] snapshot/}
   `);
 });
+
+test('should not leak websocket connections', {
+  annotation: { type: 'issue', description: 'https://github.com/microsoft/playwright/issues/33641' }
+}, async ({ runUITest }) => {
+  const { page } = await runUITest({
+    'a.test.ts': `
+      import { test, expect } from '@playwright/test';
+      test('test', async () => {});
+    `,
+  });
+
+  const [ws1] = await Promise.all([
+    page.waitForEvent('websocket'),
+    page.getByTitle('Reload').click(),
+  ]);
+
+  await Promise.all([
+    page.waitForEvent('websocket'),
+    page.getByTitle('Reload').click(),
+  ]);
+
+  await expect.poll(() => ws1.isClosed()).toBe(true);
+});