From 47dfc3578d279bec4d532a20a08a076625f373cb Mon Sep 17 00:00:00 2001
From: Nikolay Vitkov <34244704+Lightning00Blade@users.noreply.github.com>
Date: Wed, 9 Aug 2023 10:00:33 +0200
Subject: [PATCH] chore: add evaluateOnNewDocument for BiDi (#10712)

---
 .../src/common/bidi/Connection.ts             |  4 +++
 .../puppeteer-core/src/common/bidi/Page.ts    | 30 +++++++++++++++++++
 test/TestExpectations.json                    | 12 ++++++++
 3 files changed, 46 insertions(+)

diff --git a/packages/puppeteer-core/src/common/bidi/Connection.ts b/packages/puppeteer-core/src/common/bidi/Connection.ts
index 2a5265abe2c..4fcd97309f6 100644
--- a/packages/puppeteer-core/src/common/bidi/Connection.ts
+++ b/packages/puppeteer-core/src/common/bidi/Connection.ts
@@ -47,6 +47,10 @@ interface Commands {
     params: Bidi.Script.AddPreloadScriptParameters;
     returnType: Bidi.Script.AddPreloadScriptResult;
   };
+  'script.removePreloadScript': {
+    params: Bidi.Script.RemovePreloadScriptParameters;
+    returnType: Bidi.EmptyResult;
+  };
 
   'browsingContext.activate': {
     params: Bidi.BrowsingContext.ActivateParameters;
diff --git a/packages/puppeteer-core/src/common/bidi/Page.ts b/packages/puppeteer-core/src/common/bidi/Page.ts
index e3ed18b1ac0..89e54c87d71 100644
--- a/packages/puppeteer-core/src/common/bidi/Page.ts
+++ b/packages/puppeteer-core/src/common/bidi/Page.ts
@@ -22,6 +22,7 @@ import Protocol from 'devtools-protocol';
 import {
   GeolocationOptions,
   MediaFeature,
+  NewDocumentScriptEvaluation,
   Page as PageBase,
   PageEmittedEvents,
   ScreenshotOptions,
@@ -45,6 +46,7 @@ import {Tracing} from '../Tracing.js';
 import {EvaluateFunc, HandleFor} from '../types.js';
 import {
   debugError,
+  evaluationString,
   isString,
   validateDialogType,
   waitForEvent,
@@ -726,6 +728,30 @@ export class Page extends PageBase {
       context: this.mainFrame()._id,
     });
   }
+
+  override async evaluateOnNewDocument<
+    Params extends unknown[],
+    Func extends (...args: Params) => unknown = (...args: Params) => unknown,
+  >(
+    pageFunction: Func | string,
+    ...args: Params
+  ): Promise<NewDocumentScriptEvaluation> {
+    const expression = evaluationExpression(pageFunction, ...args);
+    const {result} = await this.#connection.send('script.addPreloadScript', {
+      functionDeclaration: expression,
+      // TODO: should change spec to accept browsingContext
+    });
+
+    return {identifier: result.script};
+  }
+
+  override async removeScriptToEvaluateOnNewDocument(
+    id: string
+  ): Promise<void> {
+    await this.#connection.send('script.removePreloadScript', {
+      script: id,
+    });
+  }
 }
 
 function isConsoleLogEntry(
@@ -755,3 +781,7 @@ function getStackTraceLocations(
   }
   return stackTraceLocations;
 }
+
+function evaluationExpression(fun: Function | string, ...args: unknown[]) {
+  return `() => {${evaluationString(fun, ...args)}}`;
+}
diff --git a/test/TestExpectations.json b/test/TestExpectations.json
index ba51e60ad2c..362daad7b92 100644
--- a/test/TestExpectations.json
+++ b/test/TestExpectations.json
@@ -4150,5 +4150,17 @@
     "platforms": ["darwin", "linux", "win32"],
     "parameters": ["cdp", "chrome", "headless"],
     "expectations": ["FAIL", "PASS"]
+  },
+  {
+    "testIdPattern": "[evaluation.spec] Evaluation specs Page.evaluateOnNewDocument should evaluate before anything else on the page",
+    "platforms": ["darwin", "linux", "win32"],
+    "parameters": ["webDriverBiDi"],
+    "expectations": ["PASS"]
+  },
+  {
+    "testIdPattern": "[evaluation.spec] Evaluation specs Page.evaluateOnNewDocument should work with CSP",
+    "platforms": ["darwin", "linux", "win32"],
+    "parameters": ["webDriverBiDi"],
+    "expectations": ["PASS"]
   }
 ]