fix: mimic rejection for PuppeteerUtil on early call (#9589)

puppeteer · Jan 26, 2023 · 1980de9 · 1980de9
1 parent 6edd996
commit 1980de9
Show file tree

Hide file tree

Showing 3 changed files with 35 additions and 17 deletions.
diff --git a/packages/puppeteer-core/src/common/IsolatedWorld.ts b/packages/puppeteer-core/src/common/IsolatedWorld.ts
@@ -25,7 +25,6 @@ import {Frame} from './Frame.js';
 import {FrameManager} from './FrameManager.js';
 import {MouseButton} from './Input.js';
 import {JSHandle} from './JSHandle.js';
-import {LazyArg} from './LazyArg.js';
 import {LifecycleWatcher, PuppeteerLifeCycleEvent} from './LifecycleWatcher.js';
 import {TimeoutSettings} from './TimeoutSettings.js';
 import {EvaluateFunc, HandleFor, InnerLazyParams, NodeFor} from './types.js';
@@ -35,6 +34,7 @@ import {MAIN_WORLD, PUPPETEER_WORLD} from './IsolatedWorlds.js';
 
 import type PuppeteerUtil from '../injected/injected.js';
 import type {ElementHandle} from './ElementHandle.js';
+import {LazyArg} from './LazyArg.js';
 
 /**
  * @public
@@ -96,10 +96,20 @@ export class IsolatedWorld {
   // Contains mapping from functions that should be bound to Puppeteer functions.
   #boundFunctions = new Map<string, Function>();
   #taskManager = new TaskManager();
-  #puppeteerUtil = createDeferredPromise<JSHandle<PuppeteerUtil>>();
+  #puppeteerUtil?: Promise<JSHandle<PuppeteerUtil> | undefined>;
 
   get puppeteerUtil(): Promise<JSHandle<PuppeteerUtil>> {
-    return this.#puppeteerUtil;
+    /**
+     * This is supposed to mimic what happens when evaluating Puppeteer utilities
+     * break due to navigation.
+     */
+    return (async () => {
+      const util = await this.#puppeteerUtil;
+      if (util) {
+        return util;
+      }
+      throw new Error('Execution context was destroyed!');
+    })();
   }
 
   get taskManager(): TaskManager {
@@ -151,15 +161,19 @@ export class IsolatedWorld {
 
   async #injectPuppeteerUtil(context: ExecutionContext): Promise<void> {
     try {
-      this.#puppeteerUtil.resolve(
-        (await context.evaluateHandle(
-          `(() => {
-              const module = {};
-              ${injectedSource}
-              return module.exports.default;
-            })()`
-        )) as JSHandle<PuppeteerUtil>
-      );
+      this.#puppeteerUtil = (async () => {
+        try {
+          return (await context.evaluateHandle(
+            `(() => {
+            const module = {};
+            ${injectedSource}
+            return module.exports.default;
+          })()`
+          )) as JSHandle<PuppeteerUtil>;
+        } catch {
+          return undefined;
+        }
+      })();
       this.#taskManager.rerunAll();
     } catch (error: unknown) {
       debugError(error);

diff --git a/packages/puppeteer-core/src/common/QueryHandler.ts b/packages/puppeteer-core/src/common/QueryHandler.ts
@@ -15,6 +15,7 @@
  */
 
 import PuppeteerUtil from '../injected/injected.js';
+import {assert} from '../util/assert.js';
 import {ariaHandler} from './AriaQueryHandler.js';
 import {ElementHandle} from './ElementHandle.js';
 import {Frame} from './Frame.js';
@@ -99,10 +100,12 @@ function createPuppeteerQueryHandler(
   if (handler.queryOne) {
     const queryOne = handler.queryOne;
     internalHandler.queryOne = async (element, selector) => {
+      const world = element.executionContext()._world;
+      assert(world);
       const jsHandle = await element.evaluateHandle(
         queryOne,
         selector,
-        await element.executionContext()._world!.puppeteerUtil
+        await world.puppeteerUtil
       );
       const elementHandle = jsHandle.asElement();
       if (elementHandle) {
@@ -145,6 +148,8 @@ function createPuppeteerQueryHandler(
   if (handler.queryAll) {
     const queryAll = handler.queryAll;
     internalHandler.queryAll = async (element, selector) => {
+      const world = element.executionContext()._world;
+      assert(world);
       const jsHandle = await element.evaluateHandle(
         queryAll,
         selector,

diff --git a/packages/puppeteer-core/src/common/WaitTask.ts b/packages/puppeteer-core/src/common/WaitTask.ts
@@ -170,11 +170,10 @@ export class WaitTask<T = unknown> {
   async terminate(error?: unknown): Promise<void> {
     this.#world.taskManager.delete(this);
 
-    if (this.#timeout) {
-      clearTimeout(this.#timeout);
-    }
-
     if (error && !this.#result.finished()) {
+      if (this.#timeout) {
+        clearTimeout(this.#timeout);
+      }
       this.#result.reject(error);
     }