From ff243f1af59dcbf8f3e22768491e0c6ab04080a7 Mon Sep 17 00:00:00 2001
From: Dmitry Gozman <dgozman@gmail.com>
Date: Tue, 2 Mar 2021 16:03:48 -0800
Subject: [PATCH] fix(addInitScript): make it work on new pages without
 navigations (#5675)

---
 src/server/chromium/crPage.ts               |  4 ++-
 src/server/webkit/wkPage.ts                 |  6 ++--
 test/browsercontext-expose-function.spec.ts |  3 +-
 test/page-add-init-script.spec.ts           | 33 +++++++++++++++++++++
 4 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/src/server/chromium/crPage.ts b/src/server/chromium/crPage.ts
index 3cd841bb82..c90b1bc1bf 100644
--- a/src/server/chromium/crPage.ts
+++ b/src/server/chromium/crPage.ts
@@ -421,7 +421,9 @@ class FrameSession {
             worldName: UTILITY_WORLD_NAME,
           });
           for (const binding of this._crPage._browserContext._pageBindings.values())
-            frame._evaluateExpression(binding.source, false, {}, binding.world).catch(e => {});
+            frame._evaluateExpression(binding.source, false, undefined, binding.world).catch(e => {});
+          for (const source of this._crPage._browserContext._evaluateOnNewDocumentSources)
+            frame._evaluateExpression(source, false, undefined, 'main').catch(e => {});
         }
         const isInitialEmptyPage = this._isMainFrame() && this._page.mainFrame().url() === ':';
         if (isInitialEmptyPage) {
diff --git a/src/server/webkit/wkPage.ts b/src/server/webkit/wkPage.ts
index b269a5b609..2a3169e60b 100644
--- a/src/server/webkit/wkPage.ts
+++ b/src/server/webkit/wkPage.ts
@@ -176,9 +176,9 @@ export class WKPage implements PageDelegate {
       promises.push(session.send('Page.overrideUserAgent', { value: contextOptions.userAgent }));
     if (this._page._state.mediaType || this._page._state.colorScheme)
       promises.push(WKPage._setEmulateMedia(session, this._page._state.mediaType, this._page._state.colorScheme));
-    promises.push(session.send('Page.setBootstrapScript', { source: this._calculateBootstrapScript() }));
-    for (const binding of this._browserContext._pageBindings.values())
-      promises.push(this._evaluateBindingScript(binding));
+    const bootstrapScript = this._calculateBootstrapScript();
+    promises.push(session.send('Page.setBootstrapScript', { source: bootstrapScript }));
+    this._page.frames().map(frame => frame._evaluateExpression(bootstrapScript, false, undefined, 'main').catch(e => {}));
     if (contextOptions.bypassCSP)
       promises.push(session.send('Page.setBypassCSP', { enabled: true }));
     if (this._page._state.viewportSize) {
diff --git a/test/browsercontext-expose-function.spec.ts b/test/browsercontext-expose-function.spec.ts
index 96a74a5299..48924c913f 100644
--- a/test/browsercontext-expose-function.spec.ts
+++ b/test/browsercontext-expose-function.spec.ts
@@ -66,7 +66,8 @@ it('should be callable from-inside addInitScript', async ({context, server}) =>
   await page.addInitScript('window["woof"]("page")');
   args = [];
   await page.reload();
-  expect(args).toEqual(['context', 'page']);
+  expect(args).toContain('context');
+  expect(args).toContain('page');
 });
 
 it('exposeBindingHandle should work', async ({context}) => {
diff --git a/test/page-add-init-script.spec.ts b/test/page-add-init-script.spec.ts
index ffb2296cab..44af90b510 100644
--- a/test/page-add-init-script.spec.ts
+++ b/test/page-add-init-script.spec.ts
@@ -54,6 +54,39 @@ it('should work with browser context scripts', async ({ browser, server }) => {
   await context.close();
 });
 
+it('should work without navigation, after all bindings', async ({ browser }) => {
+  const context = await browser.newContext();
+
+  let callback;
+  const promise = new Promise(f => callback = f);
+  await context.exposeFunction('woof', function(arg) {
+    callback(arg);
+  });
+
+  await context.addInitScript(() => {
+    window['woof']('hey');
+    window['temp'] = 123;
+  });
+  const page = await context.newPage();
+
+  expect(await page.evaluate(() => window['temp'])).toBe(123);
+  expect(await promise).toBe('hey');
+
+  await context.close();
+});
+
+it('should work without navigation in popup', async ({ browser }) => {
+  const context = await browser.newContext();
+  await context.addInitScript(() => window['temp'] = 123);
+  const page = await context.newPage();
+  const [popup] = await Promise.all([
+    page.waitForEvent('popup'),
+    page.evaluate(() => window['win'] = window.open()),
+  ]);
+  expect(await popup.evaluate(() => window['temp'])).toBe(123);
+  await context.close();
+});
+
 it('should work with browser context scripts with a path', async ({ browser, server }) => {
   const context = await browser.newContext();
   await context.addInitScript({ path: path.join(__dirname, 'assets/injectedfile.js') });