From 3bded358342c9b73a872d1bfa783bcac2d50eb47 Mon Sep 17 00:00:00 2001
From: Joel Einbinder <joel.einbinder@gmail.com>
Date: Thu, 13 May 2021 14:10:52 -0700
Subject: [PATCH] fix(chromium): wait for existing pages when connecting
 (#6511)

---
 src/client/electron.ts              | 14 +++++++++-----
 src/server/chromium/crBrowser.ts    | 14 +++++++++++---
 src/server/electron/electron.ts     |  4 +++-
 tests/chromium/chromium.spec.ts     |  1 -
 tests/electron/electron-app.spec.ts |  2 +-
 5 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/src/client/electron.ts b/src/client/electron.ts
index b1540e5a99..255df6027b 100644
--- a/src/client/electron.ts
+++ b/src/client/electron.ts
@@ -67,14 +67,18 @@ export class ElectronApplication extends ChannelOwner<channels.ElectronApplicati
   constructor(parent: ChannelOwner, type: string, guid: string, initializer: channels.ElectronApplicationInitializer) {
     super(parent, type, guid, initializer);
     this._context = BrowserContext.from(initializer.context);
-    this._context.on(Events.BrowserContext.Page, page => {
-      this._windows.add(page);
-      this.emit(Events.ElectronApplication.Window, page);
-      page.once(Events.Page.Close, () => this._windows.delete(page));
-    });
+    for (const page of this._context._pages)
+      this._onPage(page);
+    this._context.on(Events.BrowserContext.Page, page => this._onPage(page));
     this._channel.on('close', () => this.emit(Events.ElectronApplication.Close));
   }
 
+  _onPage(page: Page) {
+    this._windows.add(page);
+    this.emit(Events.ElectronApplication.Window, page);
+    page.once(Events.Page.Close, () => this._windows.delete(page));
+  }
+
   windows(): Page[] {
     // TODO: add ElectronPage class inherting from Page.
     return [...this._windows];
diff --git a/src/server/chromium/crBrowser.ts b/src/server/chromium/crBrowser.ts
index 269b3c0a80..877dddd80e 100644
--- a/src/server/chromium/crBrowser.ts
+++ b/src/server/chromium/crBrowser.ts
@@ -61,12 +61,16 @@ export class CRBrowser extends Browser {
       return browser;
     }
     browser._defaultContext = new CRBrowserContext(browser, undefined, options.persistent);
-
     await Promise.all([
-      session.send('Target.setAutoAttach', { autoAttach: true, waitForDebuggerOnStart: true, flatten: true }),
+      session.send('Target.setAutoAttach', { autoAttach: true, waitForDebuggerOnStart: true, flatten: true }).then(async () => {
+        // Target.setAutoAttach has a bug where it does not wait for new Targets being attached.
+        // However making a dummy call afterwards fixes this.
+        // This can be removed after https://chromium-review.googlesource.com/c/chromium/src/+/2885888 lands in stable.
+        await session.send('Target.getTargetInfo');
+      }),
       (browser._defaultContext as CRBrowserContext)._initialize(),
     ]);
-
+    await browser._waitForAllPagesToBeInitialized();
     return browser;
   }
 
@@ -104,6 +108,10 @@ export class CRBrowser extends Browser {
     return this.options.name === 'clank';
   }
 
+  async _waitForAllPagesToBeInitialized() {
+    await Promise.all([...this._crPages.values()].map(page => page.pageOrError()));
+  }
+
   _onAttachedToTarget({targetInfo, sessionId, waitingForDebugger}: Protocol.Target.attachedToTargetPayload) {
     if (targetInfo.type === 'browser')
       return;
diff --git a/src/server/electron/electron.ts b/src/server/electron/electron.ts
index 55ff02ee5a..65fa14f3b1 100644
--- a/src/server/electron/electron.ts
+++ b/src/server/electron/electron.ts
@@ -63,6 +63,8 @@ export class ElectronApplication extends SdkObject {
       // Emit application closed after context closed.
       Promise.resolve().then(() => this.emit(ElectronApplication.Events.Close));
     });
+    for (const page of this._browserContext.pages())
+      this._onPage(page);
     this._browserContext.on(BrowserContext.Events.Page, event => this._onPage(event));
     this._nodeConnection = nodeConnection;
     this._nodeSession = nodeConnection.rootSession;
@@ -77,7 +79,7 @@ export class ElectronApplication extends SdkObject {
     this._nodeSession.send('Runtime.enable', {}).catch(e => {});
   }
 
-  private async _onPage(page: Page) {
+  private _onPage(page: Page) {
     // Needs to be sync.
     const windowId = ++this._lastWindowId;
     (page as any)._browserWindowId = windowId;
diff --git a/tests/chromium/chromium.spec.ts b/tests/chromium/chromium.spec.ts
index 77e0adc22d..6da9a0930b 100644
--- a/tests/chromium/chromium.spec.ts
+++ b/tests/chromium/chromium.spec.ts
@@ -240,7 +240,6 @@ playwrightTest('should send extra headers with connect request', async ({browser
 });
 
 playwrightTest('should report all pages in an existing browser', async ({ browserType, browserOptions }, testInfo) => {
-  playwrightTest.fail();
   const port = 9339 + testInfo.workerIndex;
   const browserServer = await browserType.launch({
     ...browserOptions,
diff --git a/tests/electron/electron-app.spec.ts b/tests/electron/electron-app.spec.ts
index 6ceed9f54d..8f50292743 100644
--- a/tests/electron/electron-app.spec.ts
+++ b/tests/electron/electron-app.spec.ts
@@ -102,7 +102,7 @@ test('should return browser window', async ({ playwright }) => {
   const electronApp = await playwright._electron.launch({
     args: [path.join(__dirname, 'electron-window-app.js')],
   });
-  const page = await electronApp.waitForEvent('window');
+  const page = await electronApp.firstWindow();
   const bwHandle = await electronApp.browserWindow(page);
   expect(await bwHandle.evaluate((bw: BrowserWindow) => bw.title)).toBe('Electron');
   await electronApp.close();