feature: $wait similar to waitForFunction, but taking a selector (#303)

src/dom.ts

@@ -377,7 +377,7 @@ function normalizeSelector(selector: string): string {
 
 export type Task = (context: FrameExecutionContext) => Promise<js.JSHandle>;
 
-export function waitForFunctionTask(pageFunction: Function | string, options: types.WaitForFunctionOptions, ...args: any[]) {
+export function waitForFunctionTask(selector: string | undefined, pageFunction: Function | string, options: types.WaitForFunctionOptions, ...args: any[]) {
   const { polling = 'raf' } = options;
   if (helper.isString(polling))
     assert(polling === 'raf' || polling === 'mutation', 'Unknown polling option: ' + polling);
@@ -386,33 +386,40 @@ export function waitForFunctionTask(pageFunction: Function | string, options: ty
   else
     throw new Error('Unknown polling options: ' + polling);
   const predicateBody = helper.isString(pageFunction) ? 'return (' + pageFunction + ')' : 'return (' + pageFunction + ')(...args)';
+  if (selector !== undefined)
+    selector = normalizeSelector(selector);
 
-  return async (context: FrameExecutionContext) => context.evaluateHandle((injected: Injected, predicateBody: string, polling: types.Polling, timeout: number, ...args) => {
+  return async (context: FrameExecutionContext) => context.evaluateHandle((injected: Injected, selector: string | undefined, predicateBody: string, polling: types.Polling, timeout: number, ...args) => {
-    const predicate = new Function('...args', predicateBody);
+    const innerPredicate = new Function('...args', predicateBody);
     if (polling === 'raf')
-      return injected.pollRaf(predicate, timeout, ...args);
+      return injected.pollRaf(selector, predicate, timeout);
     if (polling === 'mutation')
-      return injected.pollMutation(predicate, timeout, ...args);
+      return injected.pollMutation(selector, predicate, timeout);
-    return injected.pollInterval(polling, predicate, timeout, ...args);
+    return injected.pollInterval(selector, polling, predicate, timeout);
-  }, await context._injected(), predicateBody, polling, options.timeout, ...args);
+
+    function predicate(element: Element | undefined): any {
+      if (selector === undefined)
+        return innerPredicate(...args);
+      return innerPredicate(element, ...args);
+    }
+  }, await context._injected(), selector, predicateBody, polling, options.timeout, ...args);
 }
 
 export function waitForSelectorTask(selector: string, visibility: types.Visibility | undefined, timeout: number): Task {
   return async (context: FrameExecutionContext) => {
     selector = normalizeSelector(selector);
-    return context.evaluateHandle((injected: Injected, selector: string, visibility: types.Visibility, timeout: number, scope?: Node) => {
+    return context.evaluateHandle((injected: Injected, selector: string, visibility: types.Visibility, timeout: number) => {
       if (visibility !== 'any')
-        return injected.pollRaf(predicate, timeout);
+        return injected.pollRaf(selector, predicate, timeout);
-      return injected.pollMutation(predicate, timeout);
+      return injected.pollMutation(selector, predicate, timeout);
 
-      function predicate(): Element | boolean {
+      function predicate(element: Element | undefined): Element | boolean {
-        const element = injected.querySelector(selector, scope || document);
         if (!element)
           return visibility === 'hidden';
         if (visibility === 'any')
           return element;
         return injected.isVisible(element) === (visibility === 'visible') ? element : false;
       }
-    }, await context._injected(), selector, visibility, timeout, undefined);
+    }, await context._injected(), selector, visibility, timeout);
   };
 }

src/frames.ts

@@ -701,9 +701,15 @@ export class Frame {
     return result.asElement() as dom.ElementHandle<Element>;
   }
 
-  waitForFunction(pageFunction: Function | string, options: types.WaitForFunctionOptions = {}, ...args: any[]): Promise<js.JSHandle> {
+  waitForFunction(pageFunction: Function | string, options?: types.WaitForFunctionOptions, ...args: any[]): Promise<js.JSHandle> {
-    options = { timeout: this._page._timeoutSettings.timeout(), ...options };
+    options = { timeout: this._page._timeoutSettings.timeout(), ...(options || {}) };
-    const task = dom.waitForFunctionTask(pageFunction, options, ...args);
+    const task = dom.waitForFunctionTask(undefined, pageFunction, options, ...args);
+    return this._scheduleRerunnableTask(task, 'main', options.timeout);
+  }
+
+  async $wait(selector: string, pageFunction: Function | string, options?: types.WaitForFunctionOptions, ...args: any[]): Promise<js.JSHandle> {
+    options = { timeout: this._page._timeoutSettings.timeout(), ...(options || {}) };
+    const task = dom.waitForFunctionTask(selector, pageFunction, options, ...args);
     return this._scheduleRerunnableTask(task, 'main', options.timeout);
   }
 

src/injected/injected.ts

@@ -30,6 +30,7 @@ function createAttributeEngine(attribute: string): SelectorEngine {
 }
 
 type ParsedSelector = { engine: SelectorEngine, selector: string }[];
+type Predicate = (element: Element | undefined) => any;
 
 class Injected {
   readonly utils: Utils;
@@ -129,23 +130,26 @@ class Injected {
     return !!(rect.top || rect.bottom || rect.width || rect.height);
   }
 
-  pollMutation(predicate: Function, timeout: number, ...args: any[]): Promise<any> {
+  pollMutation(selector: string | undefined, predicate: Predicate, timeout: number): Promise<any> {
     let timedOut = false;
     if (timeout)
       setTimeout(() => timedOut = true, timeout);
 
-    const success = predicate.apply(null, args);
+    const element = selector === undefined ? undefined : this.querySelector(selector, document);
+    const success = predicate(element);
     if (success)
       return Promise.resolve(success);
 
     let fulfill: (result?: any) => void;
     const result = new Promise(x => fulfill = x);
-    const observer = new MutationObserver(mutations => {
+    const observer = new MutationObserver(() => {
       if (timedOut) {
         observer.disconnect();
         fulfill();
+        return;
       }
-      const success = predicate.apply(null, args);
+      const element = selector === undefined ? undefined : this.querySelector(selector, document);
+      const success = predicate(element);
       if (success) {
         observer.disconnect();
         fulfill(success);
@@ -159,50 +163,53 @@ class Injected {
     return result;
   }
 
-  pollRaf(predicate: Function, timeout: number, ...args: any[]): Promise<any> {
+  pollRaf(selector: string | undefined, predicate: Predicate, timeout: number): Promise<any> {
     let timedOut = false;
     if (timeout)
       setTimeout(() => timedOut = true, timeout);
 
     let fulfill: (result?: any) => void;
     const result = new Promise(x => fulfill = x);
-    onRaf();
-    return result;
 
-    function onRaf() {
+    const onRaf = () => {
       if (timedOut) {
         fulfill();
         return;
       }
-      const success = predicate.apply(null, args);
+      const element = selector === undefined ? undefined : this.querySelector(selector, document);
+      const success = predicate(element);
       if (success)
         fulfill(success);
       else
         requestAnimationFrame(onRaf);
-    }
+    };
+
+    onRaf();
+    return result;
   }
 
-  pollInterval(pollInterval: number, predicate: Function, timeout: number, ...args: any[]): Promise<any> {
+  pollInterval(selector: string | undefined, pollInterval: number, predicate: Predicate, timeout: number): Promise<any> {
     let timedOut = false;
     if (timeout)
       setTimeout(() => timedOut = true, timeout);
 
     let fulfill: (result?: any) => void;
     const result = new Promise(x => fulfill = x);
-    onTimeout();
+    const onTimeout = () => {
-    return result;
-
-    function onTimeout() {
       if (timedOut) {
         fulfill();
         return;
       }
-      const success = predicate.apply(null, args);
+      const element = selector === undefined ? undefined : this.querySelector(selector, document);
+      const success = predicate(element);
       if (success)
         fulfill(success);
       else
         setTimeout(onTimeout, pollInterval);
-    }
+    };
+
+    onTimeout();
+    return result;
   }
 }
 

src/page.ts

@@ -445,7 +445,11 @@ export class Page extends EventEmitter {
     return this.mainFrame().waitFor(selectorOrFunctionOrTimeout, options, ...args);
   }
 
-  waitForFunction(pageFunction: Function | string, options: types.WaitForFunctionOptions, ...args: any[]): Promise<js.JSHandle> {
+  async waitForFunction(pageFunction: Function | string, options?: types.WaitForFunctionOptions, ...args: any[]): Promise<js.JSHandle> {
     return this.mainFrame().waitForFunction(pageFunction, options, ...args);
   }
+
+  async $wait(selector: string, pageFunction: Function | string, options?: types.WaitForFunctionOptions, ...args: any[]): Promise<js.JSHandle> {
+    return this.mainFrame().$wait(selector, pageFunction, options, ...args);
+  }
 }

test/page.spec.js

@@ -169,7 +169,7 @@ module.exports.describe = function({testRunner, expect, headless, playwright, FF
       expect(await page.evaluate(() => !!window.opener)).toBe(false);
       expect(await popup.evaluate(() => !!window.opener)).toBe(false);
     });
-    it.skip(WEBKIT)('should not treat navigations as new popups', async({page, server}) => {
+    it.skip(WEBKIT || FFOX)('should not treat navigations as new popups', async({page, server}) => {
       await page.goto(server.EMPTY_PAGE);
       await page.setContent('<a target=_blank rel=noopener href="/one-style.html">yo</a>');
       const [popup] = await Promise.all([

test/waittask.spec.js

@@ -205,6 +205,36 @@ module.exports.describe = function({testRunner, expect, product, playwright, FFO
     });
   });
 
+  describe('Frame.$wait', function() {
+    it('should accept arguments', async({page, server}) => {
+      await page.setContent('<div></div>');
+      const result = await page.$wait('div', (e, foo, bar) => e.nodeName + foo + bar, {}, 'foo1', 'bar2');
+      expect(await result.jsonValue()).toBe('DIVfoo1bar2');
+    });
+    it('should query selector constantly', async({page, server}) => {
+      await page.setContent('<div></div>');
+      let done = null;
+      const resultPromise = page.$wait('span', e => e).then(r => done = r);
+      expect(done).toBe(null);
+      await page.setContent('<section></section>');
+      expect(done).toBe(null);
+      await page.setContent('<span>text</span>');
+      await resultPromise;
+      expect(done).not.toBe(null);
+      expect(await done.evaluate(e => e.textContent)).toBe('text');
+    });
+    it('should be able to wait for removal', async({page}) => {
+      await page.setContent('<div></div>');
+      let done = null;
+      const resultPromise = page.$wait('div', e => !e).then(r => done = r);
+      expect(done).toBe(null);
+      await page.setContent('<section></section>');
+      await resultPromise;
+      expect(done).not.toBe(null);
+      expect(await done.jsonValue()).toBe(true);
+    });
+  });
+
   describe('Frame.waitForSelector', function() {
     const addElement = tag => document.body.appendChild(document.createElement(tag));