fix(launch): handle websocket connect exceptions (#2184)

Author: dgozman

src/server/chromium.ts

@@ -130,7 +130,9 @@ export class Chromium extends AbstractBrowserType<CRBrowser> {
   }
 
   async connect(options: ConnectOptions): Promise<CRBrowser> {
-    return await WebSocketTransport.connect(options.wsEndpoint, transport => {
+    return await WebSocketTransport.connect(options.wsEndpoint, async transport => {
+      if ((options as any).__testHookBeforeCreateBrowser)
+        await (options as any).__testHookBeforeCreateBrowser();
       return CRBrowser.connect(transport, false, new RootLogger(options.logger), options);
     });
   }

src/server/firefox.ts

@@ -148,9 +148,10 @@ export class Firefox extends AbstractBrowserType<FFBrowser> {
   }
 
   async connect(options: ConnectOptions): Promise<FFBrowser> {
-    const logger = new RootLogger(options.logger);
-    return await WebSocketTransport.connect(options.wsEndpoint, transport => {
-      return FFBrowser.connect(transport, logger, false, options.slowMo);
+    return await WebSocketTransport.connect(options.wsEndpoint, async transport => {
+      if ((options as any).__testHookBeforeCreateBrowser)
+        await (options as any).__testHookBeforeCreateBrowser();
+      return FFBrowser.connect(transport, new RootLogger(options.logger), false, options.slowMo);
     });
   }
 

src/server/webkit.ts

@@ -133,7 +133,9 @@ export class WebKit extends AbstractBrowserType<WKBrowser> {
   }
 
   async connect(options: ConnectOptions): Promise<WKBrowser> {
-    return await WebSocketTransport.connect(options.wsEndpoint, transport => {
+    return await WebSocketTransport.connect(options.wsEndpoint, async transport => {
+      if ((options as any).__testHookBeforeCreateBrowser)
+        await (options as any).__testHookBeforeCreateBrowser();
       return WKBrowser.connect(transport, new RootLogger(options.logger), options.slowMo);
     });
   }

src/transport.ts

@@ -129,7 +129,14 @@ export class WebSocketTransport implements ConnectionTransport {
   static connect<T>(url: string, onopen: (transport: ConnectionTransport) => Promise<T> | T): Promise<T> {
     const transport = new WebSocketTransport(url);
     return new Promise<T>((fulfill, reject) => {
-      transport._ws.addEventListener('open', async () => fulfill(await onopen(transport)));
+      transport._ws.addEventListener('open', async () => {
+        try {
+          fulfill(await onopen(transport));
+        } catch (e) {
+          try { transport._ws.close(); } catch (ignored) {}
+          reject(e);
+        }
+      });
       transport._ws.addEventListener('error', event => reject(new Error('WebSocket error: ' + event.message)));
     });
   }

test/launcher.spec.js

@@ -290,6 +290,15 @@ describe('browserType.connect', function() {
     await browserServer._checkLeaks();
     await browserServer.close();
   });
+  it.slow()('should handle exceptions during connect', async({browserType, defaultBrowserOptions, server}) => {
+    const browserServer = await browserType.launchServer(defaultBrowserOptions);
+    const e = new Error('Dummy');
+    const __testHookBeforeCreateBrowser = () => { throw e };
+    const error = await browserType.connect({ wsEndpoint: browserServer.wsEndpoint(), __testHookBeforeCreateBrowser }).catch(e => e);
+    await browserServer._checkLeaks();
+    await browserServer.close();
+    expect(error).toBe(e);
+  });
 });
 
 describe('browserType.launchPersistentContext', function() {