From 14f0ab7397053db5591823c716e142c684f25b44 Mon Sep 17 00:00:00 2001
From: Alex Rudenko <OrKoN@users.noreply.github.com>
Date: Wed, 30 Aug 2023 11:04:28 +0200
Subject: [PATCH] fix: apply viewport emulation to prerender targets (#10804)

---
 .../src/common/EmulationManager.ts            | 51 +++++++++++++++----
 packages/puppeteer-core/src/common/Page.ts    |  1 +
 test/TestExpectations.json                    |  6 +++
 test/src/prerender.spec.ts                    | 29 +++++++++++
 4 files changed, 78 insertions(+), 9 deletions(-)

diff --git a/packages/puppeteer-core/src/common/EmulationManager.ts b/packages/puppeteer-core/src/common/EmulationManager.ts
index 2e7ae6912f5..96b9484ae69 100644
--- a/packages/puppeteer-core/src/common/EmulationManager.ts
+++ b/packages/puppeteer-core/src/common/EmulationManager.ts
@@ -19,7 +19,7 @@ import {GeolocationOptions, MediaFeature} from '../api/Page.js';
 import {assert} from '../util/assert.js';
 import {isErrorLike} from '../util/ErrorLike.js';
 
-import {CDPSession} from './Connection.js';
+import {CDPSession, CDPSessionEmittedEvents} from './Connection.js';
 import {Viewport} from './PuppeteerViewport.js';
 
 /**
@@ -31,12 +31,24 @@ export class EmulationManager {
   #hasTouch = false;
   #javascriptEnabled = true;
 
+  #viewport?: Viewport;
+  #secondaryClients = new Set<CDPSession>();
+
   constructor(client: CDPSession) {
     this.#client = client;
   }
 
   updateClient(client: CDPSession): void {
     this.#client = client;
+    this.#secondaryClients.delete(client);
+  }
+
+  async registerSecondaryPage(client: CDPSession): Promise<void> {
+    this.#secondaryClients.add(client);
+    await this.#applyViewport(client);
+    client.once(CDPSessionEmittedEvents.Disconnected, () => {
+      return this.#secondaryClients.delete(client);
+    });
   }
 
   get javascriptEnabled(): boolean {
@@ -44,6 +56,33 @@ export class EmulationManager {
   }
 
   async emulateViewport(viewport: Viewport): Promise<boolean> {
+    this.#viewport = viewport;
+
+    await this.#applyViewport(this.#client);
+
+    const mobile = viewport.isMobile || false;
+    const hasTouch = viewport.hasTouch || false;
+    const reloadNeeded =
+      this.#emulatingMobile !== mobile || this.#hasTouch !== hasTouch;
+    this.#emulatingMobile = mobile;
+    this.#hasTouch = hasTouch;
+
+    if (!reloadNeeded) {
+      // If the page will be reloaded, no need to adjust secondary clients.
+      await Promise.all(
+        Array.from(this.#secondaryClients).map(client => {
+          return this.#applyViewport(client);
+        })
+      );
+    }
+    return reloadNeeded;
+  }
+
+  async #applyViewport(client: CDPSession): Promise<void> {
+    const viewport = this.#viewport;
+    if (!viewport) {
+      return;
+    }
     const mobile = viewport.isMobile || false;
     const width = viewport.width;
     const height = viewport.height;
@@ -55,23 +94,17 @@ export class EmulationManager {
     const hasTouch = viewport.hasTouch || false;
 
     await Promise.all([
-      this.#client.send('Emulation.setDeviceMetricsOverride', {
+      client.send('Emulation.setDeviceMetricsOverride', {
         mobile,
         width,
         height,
         deviceScaleFactor,
         screenOrientation,
       }),
-      this.#client.send('Emulation.setTouchEmulationEnabled', {
+      client.send('Emulation.setTouchEmulationEnabled', {
         enabled: hasTouch,
       }),
     ]);
-
-    const reloadNeeded =
-      this.#emulatingMobile !== mobile || this.#hasTouch !== hasTouch;
-    this.#emulatingMobile = mobile;
-    this.#hasTouch = hasTouch;
-    return reloadNeeded;
   }
 
   async emulateIdleState(overrides?: {
diff --git a/packages/puppeteer-core/src/common/Page.ts b/packages/puppeteer-core/src/common/Page.ts
index 7b362f147e2..85ea18e97ec 100644
--- a/packages/puppeteer-core/src/common/Page.ts
+++ b/packages/puppeteer-core/src/common/Page.ts
@@ -336,6 +336,7 @@ export class CDPPage extends Page {
           return;
         }
         this.#frameManager.registerSecondaryPage(session).catch(debugError);
+        this.#emulationManager.registerSecondaryPage(session).catch(debugError);
       }
     );
   }
diff --git a/test/TestExpectations.json b/test/TestExpectations.json
index 7c7a675d6df..dea439dbcc4 100644
--- a/test/TestExpectations.json
+++ b/test/TestExpectations.json
@@ -3911,6 +3911,12 @@
     "parameters": ["firefox", "webDriverBiDi"],
     "expectations": ["SKIP"]
   },
+  {
+    "testIdPattern": "[prerender.spec] Prerender with emulation can configure viewport for prerendered pages",
+    "platforms": ["darwin", "linux", "win32"],
+    "parameters": ["firefox", "webDriverBiDi"],
+    "expectations": ["SKIP"]
+  },
   {
     "testIdPattern": "[prerender.spec] Prerender with network requests can receive requests from the prerendered page",
     "platforms": ["darwin", "linux", "win32"],
diff --git a/test/src/prerender.spec.ts b/test/src/prerender.spec.ts
index 227fc0f0168..83c1191645a 100644
--- a/test/src/prerender.spec.ts
+++ b/test/src/prerender.spec.ts
@@ -129,4 +129,33 @@ describe('Prerender', function () {
       ).toBeTruthy();
     });
   });
+
+  describe('with emulation', () => {
+    it('can configure viewport for prerendered pages', async () => {
+      const {page, server} = await getTestState();
+      await page.setViewport({
+        width: 300,
+        height: 400,
+      });
+      await page.goto(server.PREFIX + '/prerender/index.html');
+      const button = await page.waitForSelector('button');
+      await button?.click();
+      const link = await page.waitForSelector('a');
+      await Promise.all([page.waitForNavigation(), link?.click()]);
+      const result = await page.evaluate(() => {
+        return {
+          width: document.documentElement.clientWidth,
+          height: document.documentElement.clientHeight,
+          dpr: window.devicePixelRatio,
+        };
+      });
+      expect({
+        width: result.width,
+        height: result.height,
+      }).toStrictEqual({
+        width: 300 * result.dpr,
+        height: 400 * result.dpr,
+      });
+    });
+  });
 });