From d54c7edeae44800dabb87e52ae96180a16a149a1 Mon Sep 17 00:00:00 2001
From: Andrey Lushnikov <aslushnikov@gmail.com>
Date: Wed, 5 Sep 2018 22:59:29 +0100
Subject: [PATCH] fix(page): Page.goto should properly handle historyAPI in
 beforeunload (#3198)

Fixes #2764.
---
 lib/NavigatorWatcher.js | 55 ++++++++++++++++++++++++-----------------
 lib/Page.js             | 22 ++++++++++++-----
 test/page.spec.js       |  2 +-
 3 files changed, 50 insertions(+), 29 deletions(-)

diff --git a/lib/NavigatorWatcher.js b/lib/NavigatorWatcher.js
index e909b5bee0a3f..92008ab8da79b 100644
--- a/lib/NavigatorWatcher.js
+++ b/lib/NavigatorWatcher.js
@@ -51,16 +51,36 @@ class NavigatorWatcher {
       helper.addEventListener(this._frameManager, FrameManager.Events.FrameDetached, this._checkLifecycleComplete.bind(this))
     ];
 
-    const lifecycleCompletePromise = new Promise(fulfill => {
-      this._lifecycleCompleteCallback = fulfill;
+    this._sameDocumentNavigationPromise = new Promise(fulfill => {
+      this._sameDocumentNavigationCompleteCallback = fulfill;
     });
-    this._navigationPromise = Promise.race([
-      this._createTimeoutPromise(),
-      lifecycleCompletePromise
-    ]).then(error => {
-      this._cleanup();
-      return error;
+
+    this._newDocumentNavigationPromise = new Promise(fulfill => {
+      this._newDocumentNavigationCompleteCallback = fulfill;
     });
+
+    this._timeoutPromise = this._createTimeoutPromise();
+  }
+
+  /**
+   * @return {!Promise<?Error>}
+   */
+  sameDocumentNavigationPromise() {
+    return this._sameDocumentNavigationPromise;
+  }
+
+  /**
+   * @return {!Promise<?Error>}
+   */
+  newDocumentNavigationPromise() {
+    return this._newDocumentNavigationPromise;
+  }
+
+  /**
+   * @return {!Promise<?Error>}
+   */
+  timeoutPromise() {
+    return this._timeoutPromise;
   }
 
   /**
@@ -74,13 +94,6 @@ class NavigatorWatcher {
         .then(() => new TimeoutError(errorMessage));
   }
 
-  /**
-   * @return {!Promise<?Error>}
-   */
-  async navigationPromise() {
-    return this._navigationPromise;
-  }
-
   /**
    * @param {!Puppeteer.Frame} frame
    */
@@ -97,7 +110,10 @@ class NavigatorWatcher {
       return;
     if (!checkLifecycle(this._frame, this._expectedLifecycle))
       return;
-    this._lifecycleCompleteCallback();
+    if (this._hasSameDocumentNavigation)
+      this._sameDocumentNavigationCompleteCallback();
+    if (this._frame._loaderId !== this._initialLoaderId)
+      this._newDocumentNavigationCompleteCallback();
 
     /**
      * @param {!Puppeteer.Frame} frame
@@ -117,13 +133,8 @@ class NavigatorWatcher {
     }
   }
 
-  cancel() {
-    this._cleanup();
-  }
-
-  _cleanup() {
+  dispose() {
     helper.removeEventListeners(this._eventListeners);
-    this._lifecycleCompleteCallback(new Error('Navigation failed'));
     clearTimeout(this._maximumTimer);
   }
 }
diff --git a/lib/Page.js b/lib/Page.js
index e53cb672eb714..b522727c89f1d 100644
--- a/lib/Page.js
+++ b/lib/Page.js
@@ -591,14 +591,18 @@ class Page extends EventEmitter {
     const mainFrame = this._frameManager.mainFrame();
     const timeout = typeof options.timeout === 'number' ? options.timeout : this._defaultNavigationTimeout;
     const watcher = new NavigatorWatcher(this._frameManager, mainFrame, timeout, options);
-    const navigationPromise = watcher.navigationPromise();
+    let ensureNewDocumentNavigation = false;
     let error = await Promise.race([
       navigate(this._client, url, referrer),
-      navigationPromise,
+      watcher.timeoutPromise(),
     ]);
-    if (!error)
-      error = await navigationPromise;
-    watcher.cancel();
+    if (!error) {
+      error = await Promise.race([
+        watcher.timeoutPromise(),
+        ensureNewDocumentNavigation ? watcher.newDocumentNavigationPromise() : watcher.sameDocumentNavigationPromise(),
+      ]);
+    }
+    watcher.dispose();
     helper.removeEventListeners(eventListeners);
     if (error)
       throw error;
@@ -614,6 +618,7 @@ class Page extends EventEmitter {
     async function navigate(client, url, referrer) {
       try {
         const response = await client.send('Page.navigate', {url, referrer});
+        ensureNewDocumentNavigation = !!response.loaderId;
         return response.errorText ? new Error(`${response.errorText} at ${url}`) : null;
       } catch (error) {
         return error;
@@ -644,7 +649,12 @@ class Page extends EventEmitter {
 
     const responses = new Map();
     const listener = helper.addEventListener(this._networkManager, NetworkManager.Events.Response, response => responses.set(response.url(), response));
-    const error = await watcher.navigationPromise();
+    const error = await Promise.race([
+      watcher.timeoutPromise(),
+      watcher.sameDocumentNavigationPromise(),
+      watcher.newDocumentNavigationPromise()
+    ]);
+    watcher.dispose();
     helper.removeEventListeners([listener]);
     if (error)
       throw error;
diff --git a/test/page.spec.js b/test/page.spec.js
index 9c1d0f9b4dfd9..b8f1421f17565 100644
--- a/test/page.spec.js
+++ b/test/page.spec.js
@@ -552,7 +552,7 @@ module.exports.addTests = function({testRunner, expect, headless}) {
       expect(response.status()).toBe(200);
       expect(response.securityDetails()).toBe(null);
     });
-    xit('should work when page calls history API in beforeunload', async({page, server}) => {
+    it('should work when page calls history API in beforeunload', async({page, server}) => {
       await page.goto(server.EMPTY_PAGE);
       await page.evaluate(() => {
         window.addEventListener('beforeunload', () => history.replaceState(null, 'initial', window.location.href), false);