Implement waitFor which survives navigation (#99)

This patch implements page.waitFor method which survives navigation.

References #89.

lib/FrameManager.js

@@ -177,7 +177,7 @@ class FrameManager extends EventEmitter {
    * @param {!Frame} frame
    * @return {!Promise<undefined>}
    */
-  async _waitForInFrame(selector, frame) {
+  async _waitForSelector(selector, frame) {
 
     function code(selector) {
       if (document.querySelector(selector))
@@ -193,7 +193,7 @@ class FrameManager extends EventEmitter {
           return;
         }
       });
-      mo.observe(document.documentElement, {
+      mo.observe(document, {
         childList: true,
         subtree: true
       });
@@ -240,6 +240,9 @@ class Frame {
     this._parentFrame = parentFrame;
     this._url = '';
     this._id = frameId;
+    /** @type {!Set<!AwaitedElement>} */
+    this._awaitedElements = new Set();
+
     this._adoptPayload(payload);
 
     /** @type {!Set<!Frame>} */
@@ -301,10 +304,15 @@ class Frame {
 
   /**
    * @param {string} selector
-   * @return {!Promise<undefined>}
+   * @return {!Promise}
    */
   async waitFor(selector) {
-    await this._frameManager._waitForInFrame(selector, this);
+    const awaitedElement = new AwaitedElement(() => this._frameManager._waitForSelector(selector, this));
+
+    this._awaitedElements.add(awaitedElement);
+    let cleanup = () => this._awaitedElements.delete(awaitedElement);
+    awaitedElement.promise.then(cleanup, cleanup);
+    return awaitedElement.promise;
   }
 
   /**
@@ -349,9 +357,13 @@ class Frame {
     };
     this._name = framePayload.name;
     this._url = framePayload.url;
+    for (let awaitedElement of this._awaitedElements)
+      awaitedElement.startWaiting();
   }
 
   _detach() {
+    for (let awaitedElement of this._awaitedElements)
+      awaitedElement.terminate(new Error('waitForSelector failed: frame got detached.'));
     this._detached = true;
     if (this._parentFrame)
       this._parentFrame._childFrames.delete(this);
@@ -360,4 +372,44 @@ class Frame {
 }
 helper.tracePublicAPI(Frame);
 
+class AwaitedElement {
+  /**
+   * @param {function():!Promise} waitInPageCallback
+   */
+  constructor(waitInPageCallback) {
+    this.promise = new Promise((resolve, reject) => {
+      this._resolve = resolve;
+      this._reject = reject;
+    });
+    this._waitInPageCallback = waitInPageCallback;
+    this._waitPromise = null;
+    this.startWaiting();
+  }
+
+  /**
+   * @param {!Error} error
+   */
+  terminate(error) {
+    this._reject(error);
+    this._waitTaskPromise = null;
+  }
+
+  startWaiting() {
+    let waitPromise = this._waitInPageCallback.call(null).then(finish.bind(this), finish.bind(this));
+    this._waitPromise = waitPromise;
+
+    /**
+     * @param {?Error} error
+     */
+    function finish(error) {
+      if (this._waitPromise !== waitPromise)
+        return;
+      if (error)
+        this._reject(error);
+      else
+        this._resolve();
+    }
+  }
+}
+
 module.exports = FrameManager;