From 5550f58284cfca2e890827298bf7442b258a9b10 Mon Sep 17 00:00:00 2001
From: Yury Semikhatsky <yurys@chromium.org>
Date: Wed, 8 Sep 2021 16:31:58 -0700
Subject: [PATCH] feat(fetch): fetch with request parameter (#8793)

---
 src/client/browserContext.ts            | 24 ++++++++++++++++++------
 src/client/page.ts                      |  6 ++++--
 tests/browsercontext-fetch.spec.ts      |  6 ++++++
 tests/page/page-request-fulfill.spec.ts | 14 ++++++++++++++
 4 files changed, 42 insertions(+), 8 deletions(-)

diff --git a/src/client/browserContext.ts b/src/client/browserContext.ts
index 5437d40968..397568db7d 100644
--- a/src/client/browserContext.ts
+++ b/src/client/browserContext.ts
@@ -28,7 +28,7 @@ import { Events } from './events';
 import { TimeoutSettings } from '../utils/timeoutSettings';
 import { Waiter } from './waiter';
 import { URLMatch, Headers, WaitForEventOptions, BrowserContextOptions, StorageState, LaunchOptions } from './types';
-import { isUnderTest, headersObjectToArray, mkdirIfNeeded, isString } from '../utils/utils';
+import { isUnderTest, headersObjectToArray, mkdirIfNeeded, isString, assert } from '../utils/utils';
 import { isSafeCloseError } from '../utils/errors';
 import * as api from '../../types/types';
 import * as structs from '../../types/structs';
@@ -216,14 +216,26 @@ export class BrowserContext extends ChannelOwner<channels.BrowserContextChannel,
     });
   }
 
-  async _fetch(url: string, options: FetchOptions = {}): Promise<network.FetchResponse> {
+  async _fetch(request: network.Request, options?: { timeout?: number }): Promise<network.FetchResponse>;
+  async _fetch(url: string, options?: FetchOptions): Promise<network.FetchResponse>;
+  async _fetch(urlOrRequest: string|network.Request, options: FetchOptions = {}): Promise<network.FetchResponse> {
     return this._wrapApiCall(async (channel: channels.BrowserContextChannel) => {
-      const postDataBuffer = isString(options.postData) ? Buffer.from(options.postData, 'utf8') : options.postData;
+      const request: network.Request | undefined = (urlOrRequest instanceof network.Request) ? urlOrRequest as network.Request : undefined;
+      assert(request || typeof urlOrRequest === 'string', 'First argument must be either URL string or Request');
+      const url = request ? request.url() : urlOrRequest as string;
+      const method = request?.method() || options.method;
+      // Cannot call allHeaders() here as the request may be paused inside route handler.
+      const headersObj = request?.headers() || options.headers;
+      const headers = headersObj ? headersObjectToArray(headersObj) : undefined;
+      let postDataBuffer = request?.postDataBuffer();
+      if (postDataBuffer === undefined)
+        postDataBuffer = (isString(options.postData) ? Buffer.from(options.postData, 'utf8') : options.postData);
+      const postData = (postDataBuffer ? postDataBuffer.toString('base64') : undefined);
       const result = await channel.fetch({
         url,
-        method: options.method,
-        headers: options.headers ? headersObjectToArray(options.headers) : undefined,
-        postData: postDataBuffer ? postDataBuffer.toString('base64') : undefined,
+        method,
+        headers,
+        postData,
         timeout: options.timeout,
       });
       if (result.error)
diff --git a/src/client/page.ts b/src/client/page.ts
index 38d1a1733f..15c2f342d6 100644
--- a/src/client/page.ts
+++ b/src/client/page.ts
@@ -438,8 +438,10 @@ export class Page extends ChannelOwner<channels.PageChannel, channels.PageInitia
     return this._mainFrame.evaluate(pageFunction, arg);
   }
 
-  async _fetch(url: string, options: FetchOptions = {}): Promise<network.FetchResponse> {
-    return await this._browserContext._fetch(url, options);
+  async _fetch(request: network.Request, options?: { timeout?: number }): Promise<network.FetchResponse>;
+  async _fetch(url: string, options?: FetchOptions): Promise<network.FetchResponse>;
+  async _fetch(urlOrRequest: string|network.Request, options: FetchOptions = {}): Promise<network.FetchResponse> {
+    return await this._browserContext._fetch(urlOrRequest as any, options);
   }
 
   async addInitScript(script: Function | string | { path?: string, content?: string }, arg?: any) {
diff --git a/tests/browsercontext-fetch.spec.ts b/tests/browsercontext-fetch.spec.ts
index ea68eba7a8..2fc88df5a4 100644
--- a/tests/browsercontext-fetch.spec.ts
+++ b/tests/browsercontext-fetch.spec.ts
@@ -632,3 +632,9 @@ it('should dispose when context closes', async function({context, server}) {
   expect(error.message).toContain('Target page, context or browser has been closed');
 });
 
+it('should throw on invalid first argument', async function({context, server}) {
+  // @ts-expect-error
+  const error = await context._fetch({}).catch(e => e);
+  expect(error.message).toContain('First argument must be either URL string or Request');
+});
+
diff --git a/tests/page/page-request-fulfill.spec.ts b/tests/page/page-request-fulfill.spec.ts
index 350010b963..5de4cf8106 100644
--- a/tests/page/page-request-fulfill.spec.ts
+++ b/tests/page/page-request-fulfill.spec.ts
@@ -224,3 +224,17 @@ it('should fulfill with fetch result and overrides', async ({page, server}) => {
   expect((await response.allHeaders()).foo).toEqual('bar');
   expect(await response.json()).toEqual({'foo': 'bar'});
 });
+
+it('should fetch original request and fulfill', async ({page, server}) => {
+  await page.route('**/*', async route => {
+    // @ts-expect-error
+    const response = await page._fetch(route.request());
+    route.fulfill({
+      // @ts-expect-error
+      response,
+    });
+  });
+  const response = await page.goto(server.PREFIX + '/title.html');
+  expect(response.status()).toBe(200);
+  expect(await page.title()).toEqual('Woof-Woof');
+});