From f342c26d953e51bf262123e2e97276eff8ee2841 Mon Sep 17 00:00:00 2001
From: Nikolay Vitkov <34244704+Lightning00Blade@users.noreply.github.com>
Date: Fri, 26 May 2023 14:50:00 +0200
Subject: [PATCH] chore: add WaitForRequest and WaitForResponse (#10257)

---
 .../puppeteer-core/src/common/bidi/Page.ts    | 57 ++++++++++++++++++-
 test/TestExpectations.json                    | 14 ++++-
 2 files changed, 67 insertions(+), 4 deletions(-)

diff --git a/packages/puppeteer-core/src/common/bidi/Page.ts b/packages/puppeteer-core/src/common/bidi/Page.ts
index 8b09cf38793..8e635aee43f 100644
--- a/packages/puppeteer-core/src/common/bidi/Page.ts
+++ b/packages/puppeteer-core/src/common/bidi/Page.ts
@@ -25,7 +25,9 @@ import {
   WaitForOptions,
 } from '../../api/Page.js';
 import {assert} from '../../util/assert.js';
+import {createDeferred} from '../../util/Deferred.js';
 import {ConsoleMessage, ConsoleMessageLocation} from '../ConsoleMessage.js';
+import {TargetCloseError} from '../Errors.js';
 import {Handler} from '../EventEmitter.js';
 import {FrameManagerEmittedEvents} from '../FrameManager.js';
 import {FrameTree} from '../FrameTree.js';
@@ -36,6 +38,8 @@ import {TimeoutSettings} from '../TimeoutSettings.js';
 import {EvaluateFunc, HandleFor} from '../types.js';
 import {
   debugError,
+  isString,
+  waitForEvent,
   waitWithTimeout,
   withSourcePuppeteerURLIfNone,
 } from '../util.js';
@@ -43,6 +47,7 @@ import {
 import {BrowsingContext, getBidiHandle} from './BrowsingContext.js';
 import {Connection} from './Connection.js';
 import {Frame} from './Frame.js';
+import {HTTPRequest} from './HTTPRequest.js';
 import {HTTPResponse} from './HTTPResponse.js';
 import {NetworkManager} from './NetworkManager.js';
 import {BidiSerializer} from './Serializer.js';
@@ -56,7 +61,7 @@ export class Page extends PageBase {
   #frameTree = new FrameTree<Frame>();
   #networkManager: NetworkManager;
   #viewport: Viewport | null = null;
-  #closed = false;
+  #closedDeferred = createDeferred<TargetCloseError>();
   #subscribedEvents = new Map<string, Handler<any>>([
     ['log.entryAdded', this.#onLogEntryAdded.bind(this)],
     [
@@ -258,10 +263,10 @@ export class Page extends PageBase {
   }
 
   override async close(): Promise<void> {
-    if (this.#closed) {
+    if (this.#closedDeferred.finished()) {
       return;
     }
-    this.#closed = true;
+    this.#closedDeferred.resolve(new TargetCloseError('Page closed!'));
     this.removeAllListeners();
     this.#networkManager.dispose();
 
@@ -445,6 +450,52 @@ export class Page extends PageBase {
 
     return buffer;
   }
+
+  override waitForRequest(
+    urlOrPredicate: string | ((req: HTTPRequest) => boolean | Promise<boolean>),
+    options: {timeout?: number} = {}
+  ): Promise<HTTPRequest> {
+    const {timeout = this.#timeoutSettings.timeout()} = options;
+    return waitForEvent(
+      this.#networkManager,
+      NetworkManagerEmittedEvents.Request,
+      async request => {
+        if (isString(urlOrPredicate)) {
+          return urlOrPredicate === request.url();
+        }
+        if (typeof urlOrPredicate === 'function') {
+          return !!(await urlOrPredicate(request));
+        }
+        return false;
+      },
+      timeout,
+      this.#closedDeferred.valueOrThrow()
+    );
+  }
+
+  override waitForResponse(
+    urlOrPredicate:
+      | string
+      | ((res: HTTPResponse) => boolean | Promise<boolean>),
+    options: {timeout?: number} = {}
+  ): Promise<HTTPResponse> {
+    const {timeout = this.#timeoutSettings.timeout()} = options;
+    return waitForEvent(
+      this.#networkManager,
+      NetworkManagerEmittedEvents.Response,
+      async response => {
+        if (isString(urlOrPredicate)) {
+          return urlOrPredicate === response.url();
+        }
+        if (typeof urlOrPredicate === 'function') {
+          return !!(await urlOrPredicate(response));
+        }
+        return false;
+      },
+      timeout,
+      this.#closedDeferred.valueOrThrow()
+    );
+  }
 }
 
 function isConsoleLogEntry(
diff --git a/test/TestExpectations.json b/test/TestExpectations.json
index c627d0bb21c..817aafeb0a0 100644
--- a/test/TestExpectations.json
+++ b/test/TestExpectations.json
@@ -129,7 +129,7 @@
     "testIdPattern": "[network.spec] network Response.text *",
     "platforms": ["darwin", "linux", "win32"],
     "parameters": ["webDriverBiDi"],
-    "expectations": ["FAIL"]
+    "expectations": ["FAIL", "TIMEOUT"]
   },
   {
     "testIdPattern": "[NetworkManager.spec] NetworkManager *",
@@ -167,6 +167,18 @@
     "parameters": ["webDriverBiDi"],
     "expectations": ["PASS"]
   },
+  {
+    "testIdPattern": "[page.spec] Page Page.waitForRequest *",
+    "platforms": ["darwin", "linux", "win32"],
+    "parameters": ["webDriverBiDi"],
+    "expectations": ["PASS"]
+  },
+  {
+    "testIdPattern": "[page.spec] Page Page.waitForResponse *",
+    "platforms": ["darwin", "linux", "win32"],
+    "parameters": ["webDriverBiDi"],
+    "expectations": ["PASS"]
+  },
   {
     "testIdPattern": "[queryhandler.spec] *",
     "platforms": ["darwin", "linux", "win32"],