From b97bddf8e5750d20c6ba82392eebe2a3fd2dd218 Mon Sep 17 00:00:00 2001
From: Andrey Lushnikov <aslushnikov@gmail.com>
Date: Tue, 18 Sep 2018 00:15:50 +0100
Subject: [PATCH] refactor: unify response tracking in page.goto and
 waitForNavigation (#3258)

This patch unifies logic in response trackign in page.goto and
page.waitForNavigation.

As a drive-by, we now make sure that we return the right response
for the right frame. This will come handy for future frame navigation
API.

References #2918
---
 lib/NavigatorWatcher.js | 26 ++++++++++++++++++++++++--
 lib/Page.js             | 23 ++++-------------------
 2 files changed, 28 insertions(+), 21 deletions(-)

diff --git a/lib/NavigatorWatcher.js b/lib/NavigatorWatcher.js
index e0ebf6c1515..b11c2c06da7 100644
--- a/lib/NavigatorWatcher.js
+++ b/lib/NavigatorWatcher.js
@@ -16,6 +16,7 @@
 
 const {helper, assert} = require('./helper');
 const {FrameManager} = require('./FrameManager');
+const {NetworkManager} = require('./NetworkManager');
 const {TimeoutError} = require('./Errors');
 const {Connection} = require('./Connection');
 
@@ -23,11 +24,12 @@ class NavigatorWatcher {
   /**
    * @param {!Puppeteer.CDPSession} client
    * @param {!FrameManager} frameManager
+   * @param {!NetworkManager} networkManager
    * @param {!Puppeteer.Frame} frame
    * @param {number} timeout
    * @param {!Object=} options
    */
-  constructor(client, frameManager, frame, timeout, options = {}) {
+  constructor(client, frameManager, networkManager, frame, timeout, options = {}) {
     assert(options.networkIdleTimeout === undefined, 'ERROR: networkIdleTimeout option is no longer supported.');
     assert(options.networkIdleInflight === undefined, 'ERROR: networkIdleInflight option is no longer supported.');
     assert(options.waitUntil !== 'networkidle', 'ERROR: "networkidle" option is no longer supported. Use "networkidle2" instead');
@@ -43,15 +45,19 @@ class NavigatorWatcher {
     });
 
     this._frameManager = frameManager;
+    this._networkManager = networkManager;
     this._frame = frame;
     this._initialLoaderId = frame._loaderId;
     this._timeout = timeout;
+    /** @type {?Puppeteer.Request} */
+    this._navigationRequest = null;
     this._hasSameDocumentNavigation = false;
     this._eventListeners = [
       helper.addEventListener(Connection.fromSession(client), Connection.Events.Disconnected, () => this._terminate(new Error('Navigation failed because browser has disconnected!'))),
       helper.addEventListener(this._frameManager, FrameManager.Events.LifecycleEvent, this._checkLifecycleComplete.bind(this)),
       helper.addEventListener(this._frameManager, FrameManager.Events.FrameNavigatedWithinDocument, this._navigatedWithinDocument.bind(this)),
-      helper.addEventListener(this._frameManager, FrameManager.Events.FrameDetached, this._checkLifecycleComplete.bind(this))
+      helper.addEventListener(this._frameManager, FrameManager.Events.FrameDetached, this._checkLifecycleComplete.bind(this)),
+      helper.addEventListener(this._networkManager, NetworkManager.Events.Request, this._onRequest.bind(this)),
     ];
 
     this._sameDocumentNavigationPromise = new Promise(fulfill => {
@@ -68,6 +74,22 @@ class NavigatorWatcher {
     });
   }
 
+  /**
+   * @param {!Puppeteer.Request} request
+   */
+  _onRequest(request) {
+    if (request.frame() !== this._frame || !request.isNavigationRequest())
+      return;
+    this._navigationRequest = request;
+  }
+
+  /**
+   * @return {?Puppeteer.Response}
+   */
+  navigationResponse() {
+    return this._navigationRequest ? this._navigationRequest.response() : null;
+  }
+
   /**
    * @param {!Error} error
    */
diff --git a/lib/Page.js b/lib/Page.js
index b73fde753b2..bdc59ba3b82 100644
--- a/lib/Page.js
+++ b/lib/Page.js
@@ -579,18 +579,9 @@ class Page extends EventEmitter {
   async goto(url, options = {}) {
     const referrer = typeof options.referer === 'string' ? options.referer : this._networkManager.extraHTTPHeaders()['referer'];
 
-    /** @type {Map<string, !Puppeteer.Request>} */
-    const requests = new Map();
-    const eventListeners = [
-      helper.addEventListener(this._networkManager, NetworkManager.Events.Request, request => {
-        if (!requests.get(request.url()))
-          requests.set(request.url(), request);
-      })
-    ];
-
     const mainFrame = this._frameManager.mainFrame();
     const timeout = typeof options.timeout === 'number' ? options.timeout : this._defaultNavigationTimeout;
-    const watcher = new NavigatorWatcher(this._client, this._frameManager, mainFrame, timeout, options);
+    const watcher = new NavigatorWatcher(this._client, this._frameManager, this._networkManager, mainFrame, timeout, options);
     let ensureNewDocumentNavigation = false;
     let error = await Promise.race([
       navigate(this._client, url, referrer),
@@ -603,11 +594,9 @@ class Page extends EventEmitter {
       ]);
     }
     watcher.dispose();
-    helper.removeEventListeners(eventListeners);
     if (error)
       throw error;
-    const request = requests.get(mainFrame._navigationURL);
-    return request ? request.response() : null;
+    return watcher.navigationResponse();
 
     /**
      * @param {!Puppeteer.CDPSession} client
@@ -645,20 +634,16 @@ class Page extends EventEmitter {
   async waitForNavigation(options = {}) {
     const mainFrame = this._frameManager.mainFrame();
     const timeout = typeof options.timeout === 'number' ? options.timeout : this._defaultNavigationTimeout;
-    const watcher = new NavigatorWatcher(this._client, this._frameManager, mainFrame, timeout, options);
-
-    const responses = new Map();
-    const listener = helper.addEventListener(this._networkManager, NetworkManager.Events.Response, response => responses.set(response.url(), response));
+    const watcher = new NavigatorWatcher(this._client, this._frameManager, this._networkManager, mainFrame, timeout, options);
     const error = await Promise.race([
       watcher.timeoutOrTerminationPromise(),
       watcher.sameDocumentNavigationPromise(),
       watcher.newDocumentNavigationPromise()
     ]);
     watcher.dispose();
-    helper.removeEventListeners([listener]);
     if (error)
       throw error;
-    return responses.get(this.mainFrame().url()) || null;
+    return watcher.navigationResponse();
   }
 
   /**