chore: move waitFor methods from DOMWorld to Frame (#87)

This almost removes the DOMWorld, so we can unify them across the browsers.

src/chromium/DOMWorld.ts

@@ -17,30 +17,23 @@
 
 import { ExecutionContext } from './ExecutionContext';
 import { Frame } from './Frame';
-import { ElementHandle, JSHandle } from './JSHandle';
+import { JSHandle } from './JSHandle';
-import { TimeoutSettings } from '../TimeoutSettings';
+import { WaitTask, WaitTaskParams } from '../waitTask';
-import { WaitTask, WaitTaskParams, waitForSelectorOrXPath } from '../waitTask';
 
 export class DOMWorld {
   private _frame: Frame;
-  private _timeoutSettings: TimeoutSettings;
   private _contextPromise: Promise<ExecutionContext>;
   private _contextResolveCallback: ((c: ExecutionContext) => void) | null;
-  private _context: ExecutionContext | null;
+  _context: ExecutionContext | null;
   _waitTasks = new Set<WaitTask<JSHandle>>();
   private _detached = false;
 
-  constructor(frame: Frame, timeoutSettings: TimeoutSettings) {
+  constructor(frame: Frame) {
     this._frame = frame;
-    this._timeoutSettings = timeoutSettings;
     this._contextPromise;
     this._setContext(null);
   }
 
-  frame(): Frame {
-    return this._frame;
-  }
-
   _setContext(context: ExecutionContext | null) {
     this._context = context;
     if (context) {
@@ -55,10 +48,6 @@ export class DOMWorld {
     }
   }
 
-  _hasContext(): boolean {
-    return !this._contextResolveCallback;
-  }
-
   _detach() {
     this._detached = true;
     for (const waitTask of this._waitTasks)
@@ -71,42 +60,7 @@ export class DOMWorld {
     return this._contextPromise;
   }
 
-  async waitForSelector(selector: string, options: { visible?: boolean; hidden?: boolean; timeout?: number; } | undefined): Promise<ElementHandle | null> {
+  scheduleWaitTask(params: WaitTaskParams): Promise<JSHandle> {
-    const params = waitForSelectorOrXPath(selector, false /* isXPath */, { timeout: this._timeoutSettings.timeout(), ...options });
-    const handle = await this._scheduleWaitTask(params);
-    if (!handle.asElement()) {
-      await handle.dispose();
-      return null;
-    }
-    return handle.asElement();
-  }
-
-  async waitForXPath(xpath: string, options: { visible?: boolean, hidden?: boolean, timeout?: number } = {}): Promise<ElementHandle | null> {
-    const params = waitForSelectorOrXPath(xpath, true /* isXPath */, { timeout: this._timeoutSettings.timeout(), ...options });
-    const handle = await this._scheduleWaitTask(params);
-    if (!handle.asElement()) {
-      await handle.dispose();
-      return null;
-    }
-    return handle.asElement();
-  }
-
-  waitForFunction(pageFunction: Function | string, options: { polling?: string | number; timeout?: number; } = {}, ...args): Promise<JSHandle> {
-    const {
-      polling = 'raf',
-      timeout = this._timeoutSettings.timeout(),
-    } = options;
-    const params: WaitTaskParams = {
-      predicateBody: pageFunction,
-      title: 'function',
-      polling,
-      timeout,
-      args
-    };
-    return this._scheduleWaitTask(params);
-  }
-
-  private _scheduleWaitTask(params: WaitTaskParams): Promise<JSHandle> {
     const task = new WaitTask(params, () => this._waitTasks.delete(task));
     this._waitTasks.add(task);
     if (this._context)

src/chromium/ExecutionContext.ts

@@ -16,7 +16,6 @@
  */
 
 import { CDPSession } from './Connection';
-import { DOMWorld } from './DOMWorld';
 import { Frame } from './Frame';
 import { assert, helper } from '../helper';
 import { valueFromRemoteObject, getExceptionMessage } from './protocolHelper';
@@ -32,19 +31,19 @@ const SOURCE_URL_REGEX = /^[\040\t]*\/\/[@#] sourceURL=\s*(\S*?)\s*$/m;
 
 export class ExecutionContext implements types.EvaluationContext<JSHandle> {
   _client: CDPSession;
-  _world: DOMWorld;
+  private _frame: Frame;
   private _injectedPromise: Promise<JSHandle> | null = null;
   private _documentPromise: Promise<ElementHandle> | null = null;
   private _contextId: number;
 
-  constructor(client: CDPSession, contextPayload: Protocol.Runtime.ExecutionContextDescription, world: DOMWorld | null) {
+  constructor(client: CDPSession, contextPayload: Protocol.Runtime.ExecutionContextDescription, frame: Frame | null) {
     this._client = client;
-    this._world = world;
+    this._frame = frame;
     this._contextId = contextPayload.id;
   }
 
   frame(): Frame | null {
-    return this._world ? this._world.frame() : null;
+    return this._frame;
   }
 
   evaluate: types.Evaluate<JSHandle> = (pageFunction, ...args) => {
@@ -154,7 +153,7 @@ export class ExecutionContext implements types.EvaluationContext<JSHandle> {
 
   async _adoptElementHandle(elementHandle: ElementHandle): Promise<ElementHandle> {
     assert(elementHandle.executionContext() !== this, 'Cannot adopt handle that already belongs to this execution context');
-    assert(this._world, 'Cannot adopt handle without DOMWorld');
+    assert(this._frame, 'Cannot adopt handle without a Frame');
     const nodeInfo = await this._client.send('DOM.describeNode', {
       objectId: elementHandle._remoteObject.objectId,
     });

src/chromium/Frame.ts

@@ -27,6 +27,7 @@ import { ElementHandle, JSHandle } from './JSHandle';
 import { Response } from './NetworkManager';
 import { Protocol } from './protocol';
 import { LifecycleWatcher } from './LifecycleWatcher';
+import { waitForSelectorOrXPath, WaitTaskParams } from '../waitTask';
 
 const readFileAsync = helper.promisify(fs.readFile);
 
@@ -51,8 +52,8 @@ export class Frame {
     this._parentFrame = parentFrame;
     this._id = frameId;
 
-    this._mainWorld = new DOMWorld(this, frameManager._timeoutSettings);
+    this._mainWorld = new DOMWorld(this);
-    this._secondaryWorld = new DOMWorld(this, frameManager._timeoutSettings);
+    this._secondaryWorld = new DOMWorld(this);
 
     if (this._parentFrame)
       this._parentFrame._childFrames.add(this);
@@ -383,11 +384,13 @@ export class Frame {
       visible?: boolean;
       hidden?: boolean;
       timeout?: number; } | undefined): Promise<ElementHandle | null> {
-    const handle = await this._secondaryWorld.waitForSelector(selector, options);
+    const params = waitForSelectorOrXPath(selector, false /* isXPath */, { timeout: this._frameManager._timeoutSettings.timeout(), ...options });
-    if (!handle)
+    const handle = await this._secondaryWorld.scheduleWaitTask(params);
-      return null;
+    let result = null;
+    if (handle.asElement()) {
       const mainExecutionContext = await this._mainWorld.executionContext();
-    const result = await mainExecutionContext._adoptElementHandle(handle);
+      result = await mainExecutionContext._adoptElementHandle(handle.asElement());
+    }
     await handle.dispose();
     return result;
   }
@@ -396,11 +399,13 @@ export class Frame {
       visible?: boolean;
       hidden?: boolean;
       timeout?: number; } | undefined): Promise<ElementHandle | null> {
-    const handle = await this._secondaryWorld.waitForXPath(xpath, options);
+    const params = waitForSelectorOrXPath(xpath, true /* isXPath */, { timeout: this._frameManager._timeoutSettings.timeout(), ...options });
-    if (!handle)
+    const handle = await this._secondaryWorld.scheduleWaitTask(params);
-      return null;
+    let result = null;
+    if (handle.asElement()) {
       const mainExecutionContext = await this._mainWorld.executionContext();
-    const result = await mainExecutionContext._adoptElementHandle(handle);
+      result = await mainExecutionContext._adoptElementHandle(handle.asElement());
+    }
     await handle.dispose();
     return result;
   }
@@ -409,7 +414,18 @@ export class Frame {
     pageFunction: Function | string,
     options: { polling?: string | number; timeout?: number; } = {},
     ...args): Promise<JSHandle> {
-    return this._mainWorld.waitForFunction(pageFunction, options, ...args);
+    const {
+      polling = 'raf',
+      timeout = this._frameManager._timeoutSettings.timeout(),
+    } = options;
+    const params: WaitTaskParams = {
+      predicateBody: pageFunction,
+      title: 'function',
+      polling,
+      timeout,
+      args
+    };
+    return this._mainWorld.scheduleWaitTask(params);
   }
 
   async title(): Promise<string> {

src/chromium/FrameManager.ts

@@ -25,6 +25,7 @@ import { LifecycleWatcher } from './LifecycleWatcher';
 import { NetworkManager, Response } from './NetworkManager';
 import { Page } from './Page';
 import { Protocol } from './protocol';
+import { DOMWorld } from './DOMWorld';
 
 const UTILITY_WORLD_NAME = '__playwright_utility_world__';
 
@@ -258,11 +259,11 @@ export class FrameManager extends EventEmitter {
   _onExecutionContextCreated(contextPayload) {
     const frameId = contextPayload.auxData ? contextPayload.auxData.frameId : null;
     const frame = this._frames.get(frameId) || null;
-    let world = null;
+    let world: DOMWorld | null = null;
     if (frame) {
       if (contextPayload.auxData && !!contextPayload.auxData['isDefault']) {
         world = frame._mainWorld;
-      } else if (contextPayload.name === UTILITY_WORLD_NAME && !frame._secondaryWorld._hasContext()) {
+      } else if (contextPayload.name === UTILITY_WORLD_NAME && !frame._secondaryWorld._context) {
         // In case of multiple sessions to the same target, there's a race between
         // connections so we might end up creating multiple isolated worlds.
         // We can use either.
@@ -271,7 +272,7 @@ export class FrameManager extends EventEmitter {
     }
     if (contextPayload.auxData && contextPayload.auxData['type'] === 'isolated')
       this._isolatedWorlds.add(contextPayload.name);
-    const context: ExecutionContext = new ExecutionContext(this._client, contextPayload, world);
+    const context: ExecutionContext = new ExecutionContext(this._client, contextPayload, frame);
     if (world)
       world._setContext(context);
     this._contextIdToContext.set(contextPayload.id, context);
@@ -282,16 +283,18 @@ export class FrameManager extends EventEmitter {
     if (!context)
       return;
     this._contextIdToContext.delete(executionContextId);
-    if (context._world)
+    const frame = context.frame();
-      context._world._setContext(null);
+    if (frame) {
+      if (frame._mainWorld._context === context)
+        frame._mainWorld._setContext(null);
+      if (frame._secondaryWorld._context === context)
+        frame._secondaryWorld._setContext(null);
+    }
   }
 
   _onExecutionContextsCleared() {
-    for (const context of this._contextIdToContext.values()) {
+    for (const contextId of Array.from(this._contextIdToContext.keys()))
-      if (context._world)
+      this._onExecutionContextDestroyed(contextId);
-        context._world._setContext(null);
-    }
-    this._contextIdToContext.clear();
   }
 
   executionContextById(contextId: number): ExecutionContext {

src/firefox/DOMWorld.ts

@@ -15,9 +15,9 @@
  * limitations under the License.
  */
 
-import {ElementHandle, JSHandle} from './JSHandle';
+import { JSHandle } from './JSHandle';
 import { ExecutionContext } from './ExecutionContext';
-import { WaitTaskParams, WaitTask, waitForSelectorOrXPath } from '../waitTask';
+import { WaitTaskParams, WaitTask } from '../waitTask';
 
 export class DOMWorld {
   _frame: any;
@@ -69,42 +69,7 @@ export class DOMWorld {
     return this._contextPromise;
   }
 
-  async waitForSelector(selector: string, options: { visible?: boolean; hidden?: boolean; timeout?: number; } | undefined): Promise<ElementHandle | null> {
+  scheduleWaitTask(params: WaitTaskParams): Promise<JSHandle> {
-    const params = waitForSelectorOrXPath(selector, false /* isXPath */, { timeout: this._timeoutSettings.timeout(), ...options });
-    const handle = await this._scheduleWaitTask(params);
-    if (!handle.asElement()) {
-      await handle.dispose();
-      return null;
-    }
-    return handle.asElement();
-  }
-
-  async waitForXPath(xpath: string, options: { visible?: boolean, hidden?: boolean, timeout?: number } = {}): Promise<ElementHandle | null> {
-    const params = waitForSelectorOrXPath(xpath, true /* isXPath */, { timeout: this._timeoutSettings.timeout(), ...options });
-    const handle = await this._scheduleWaitTask(params);
-    if (!handle.asElement()) {
-      await handle.dispose();
-      return null;
-    }
-    return handle.asElement();
-  }
-
-  waitForFunction(pageFunction: Function | string, options: { polling?: string | number; timeout?: number; } | undefined = {}, ...args): Promise<JSHandle> {
-    const {
-      polling = 'raf',
-      timeout = this._timeoutSettings.timeout(),
-    } = options;
-    const params: WaitTaskParams = {
-      predicateBody: pageFunction,
-      title: 'function',
-      polling,
-      timeout,
-      args
-    };
-    return this._scheduleWaitTask(params);
-  }
-
-  private _scheduleWaitTask(params: WaitTaskParams): Promise<JSHandle> {
     const task = new WaitTask(params, () => this._waitTasks.delete(task));
     this._waitTasks.add(task);
     if (this._context)

src/firefox/FrameManager.ts

@@ -29,6 +29,7 @@ import { TimeoutSettings } from '../TimeoutSettings';
 import { NetworkManager } from './NetworkManager';
 import { MultiClickOptions, ClickOptions, SelectOption } from '../input';
 import * as types from '../types';
+import { waitForSelectorOrXPath, WaitTaskParams } from '../waitTask';
 
 const readFileAsync = helper.promisify(fs.readFile);
 
@@ -375,16 +376,48 @@ export class Frame {
     return Promise.reject(new Error('Unsupported target type: ' + (typeof selectorOrFunctionOrTimeout)));
   }
 
-  waitForFunction(pageFunction: Function | string, options: { polling?: string | number; timeout?: number; } | undefined = {}, ...args): Promise<JSHandle> {
+  waitForFunction(
-    return this._mainWorld.waitForFunction(pageFunction, options, ...args);
+    pageFunction: Function | string,
+    options: { polling?: string | number; timeout?: number; } = {},
+    ...args): Promise<JSHandle> {
+    const {
+      polling = 'raf',
+      timeout = this._frameManager._timeoutSettings.timeout(),
+    } = options;
+    const params: WaitTaskParams = {
+      predicateBody: pageFunction,
+      title: 'function',
+      polling,
+      timeout,
+      args
+    };
+    return this._mainWorld.scheduleWaitTask(params);
   }
 
-  waitForSelector(selector: string, options: { timeout?: number; visible?: boolean; hidden?: boolean; } | undefined): Promise<ElementHandle> {
+  async waitForSelector(selector: string, options: {
-    return this._mainWorld.waitForSelector(selector, options);
+      visible?: boolean;
+      hidden?: boolean;
+      timeout?: number; } | undefined): Promise<ElementHandle | null> {
+    const params = waitForSelectorOrXPath(selector, false /* isXPath */, { timeout: this._frameManager._timeoutSettings.timeout(), ...options });
+    const handle = await this._mainWorld.scheduleWaitTask(params);
+    if (!handle.asElement()) {
+      await handle.dispose();
+      return null;
+    }
+    return handle.asElement();
   }
 
-  waitForXPath(xpath: string, options: { timeout?: number; visible?: boolean; hidden?: boolean; } | undefined): Promise<ElementHandle> {
+  async waitForXPath(xpath: string, options: {
-    return this._mainWorld.waitForXPath(xpath, options);
+      visible?: boolean;
+      hidden?: boolean;
+      timeout?: number; } | undefined): Promise<ElementHandle | null> {
+    const params = waitForSelectorOrXPath(xpath, true /* isXPath */, { timeout: this._frameManager._timeoutSettings.timeout(), ...options });
+    const handle = await this._mainWorld.scheduleWaitTask(params);
+    if (!handle.asElement()) {
+      await handle.dispose();
+      return null;
+    }
+    return handle.asElement();
   }
 
   async content(): Promise<string> {