From de5a05d93584b5287fa0b4a436a42ffc093f3824 Mon Sep 17 00:00:00 2001
From: Joel Arvidsson <joel@oblador.se>
Date: Thu, 11 Jun 2020 10:36:42 +0200
Subject: [PATCH] Listen to requests until selector is present (#243)

---
 .../browser/src/await-selector-present.js     |  21 +++
 packages/browser/src/index.js                 |   2 +
 .../src/create-chrome-target.js               | 125 ++++++++++--------
 3 files changed, 95 insertions(+), 53 deletions(-)
 create mode 100644 packages/browser/src/await-selector-present.js

diff --git a/packages/browser/src/await-selector-present.js b/packages/browser/src/await-selector-present.js
new file mode 100644
index 00000000..25f09cf8
--- /dev/null
+++ b/packages/browser/src/await-selector-present.js
@@ -0,0 +1,21 @@
+const awaitSelectorPresent = (window, selector, timeout = 10000) =>
+  new Promise((resolve, reject) => {
+    let resolutionTimer;
+    const rejectionTimer = setTimeout(() => {
+      clearTimeout(resolutionTimer);
+      reject(new Error(`Timeout after ${timeout}ms`));
+    }, timeout);
+
+    const waitForSelector = () => {
+      if (window.document.querySelector(selector)) {
+        clearTimeout(rejectionTimer);
+        resolve();
+      } else {
+        resolutionTimer = setTimeout(waitForSelector, 100);
+      }
+    };
+
+    waitForSelector();
+  });
+
+module.exports = awaitSelectorPresent;
diff --git a/packages/browser/src/index.js b/packages/browser/src/index.js
index 18a865b9..ac8d201e 100644
--- a/packages/browser/src/index.js
+++ b/packages/browser/src/index.js
@@ -1,5 +1,6 @@
 const addLokiSessionMarker = require('./add-loki-session-marker');
 const awaitLokiReady = require('./await-loki-ready');
+const awaitSelectorPresent = require('./await-selector-present');
 const createStorybookConfigurator = require('./configure-storybook');
 const disableAnimations = require('./disable-animations');
 const disableInputCaret = require('./disable-input-caret');
@@ -10,6 +11,7 @@ const getStories = require('./get-stories');
 module.exports = {
   addLokiSessionMarker,
   awaitLokiReady,
+  awaitSelectorPresent,
   createStorybookConfigurator,
   disableAnimations,
   disableInputCaret,
diff --git a/packages/target-chrome-core/src/create-chrome-target.js b/packages/target-chrome-core/src/create-chrome-target.js
index 15cc449a..ca8e0098 100644
--- a/packages/target-chrome-core/src/create-chrome-target.js
+++ b/packages/target-chrome-core/src/create-chrome-target.js
@@ -6,6 +6,7 @@ const {
   getSelectorBoxSize,
   getStories,
   awaitLokiReady,
+  awaitSelectorPresent,
   addLokiSessionMarker,
 } = require('@loki/browser');
 const {
@@ -71,64 +72,70 @@ function createChromeTarget(
       await Emulation.setEmulatedMedia({ media: options.media });
     }
 
-    const awaitRequestsFinished = () =>
-      new Promise(async (resolve, reject) => {
-        const pendingRequestURLMap = {};
-        const failedURLs = [];
-        let pageLoaded = false;
-        let stabilizationTimer = null;
-
-        const maybeFulfillPromise = () => {
-          if (pageLoaded && Object.keys(pendingRequestURLMap).length === 0) {
-            if (failedURLs.length !== 0) {
-              reject(new FetchingURLsError(failedURLs));
-            } else {
-              // In some cases such as fonts further requests will only happen after the page has been fully rendered
-              if (stabilizationTimer) {
-                clearTimeout(stabilizationTimer);
-              }
-              stabilizationTimer = setTimeout(
-                resolve,
-                REQUEST_STABILIZATION_TIMEOUT
-              );
-            }
-          }
-        };
+    const pendingRequestURLMap = {};
+    const failedURLs = [];
+    let stabilizationTimer = null;
+    let requestsFinishedAwaiter;
 
-        const requestEnded = requestId => {
-          delete pendingRequestURLMap[requestId];
-          maybeFulfillPromise();
-        };
-
-        const requestFailed = requestId => {
-          const failedURL = pendingRequestURLMap[requestId];
-          if (!fetchFailIgnore || !fetchFailIgnore.test(failedURL)) {
-            failedURLs.push(failedURL);
-          }
-          requestEnded(requestId);
-        };
+    const maybeFulfillPromise = () => {
+      if (!requestsFinishedAwaiter) {
+        return;
+      }
+      const { reject, resolve } = requestsFinishedAwaiter;
 
-        Network.requestWillBeSent(({ requestId, request }) => {
+      if (Object.keys(pendingRequestURLMap).length === 0) {
+        if (failedURLs.length !== 0) {
+          reject(new FetchingURLsError(failedURLs));
+        } else {
+          // In some cases such as fonts further requests will only happen after the page has been fully rendered
           if (stabilizationTimer) {
             clearTimeout(stabilizationTimer);
           }
-          pendingRequestURLMap[requestId] = request.url;
-        });
-
-        Network.responseReceived(({ requestId, response }) => {
-          if (response.status >= 400) {
-            requestFailed(requestId);
-          } else {
-            requestEnded(requestId);
-          }
-        });
+          stabilizationTimer = setTimeout(
+            resolve,
+            REQUEST_STABILIZATION_TIMEOUT
+          );
+        }
+      }
+    };
+
+    const startObservingRequests = () => {
+      const requestEnded = requestId => {
+        delete pendingRequestURLMap[requestId];
+        maybeFulfillPromise();
+      };
+
+      const requestFailed = requestId => {
+        const failedURL = pendingRequestURLMap[requestId];
+        if (!fetchFailIgnore || !fetchFailIgnore.test(failedURL)) {
+          failedURLs.push(failedURL);
+        }
+        requestEnded(requestId);
+      };
+
+      Network.requestWillBeSent(({ requestId, request }) => {
+        if (stabilizationTimer) {
+          clearTimeout(stabilizationTimer);
+        }
+        pendingRequestURLMap[requestId] = request.url;
+      });
 
-        Network.loadingFailed(({ requestId }) => {
+      Network.responseReceived(({ requestId, response }) => {
+        if (response.status >= 400) {
           requestFailed(requestId);
-        });
+        } else {
+          requestEnded(requestId);
+        }
+      });
 
-        await Page.loadEventFired();
-        pageLoaded = true;
+      Network.loadingFailed(({ requestId }) => {
+        requestFailed(requestId);
+      });
+    };
+
+    const awaitRequestsFinished = () =>
+      new Promise((resolve, reject) => {
+        requestsFinishedAwaiter = { resolve, reject };
         maybeFulfillPromise();
       });
 
@@ -159,7 +166,7 @@ function createChromeTarget(
 
     client.executeFunctionWithWindow = executeFunctionWithWindow;
 
-    client.loadUrl = async url => {
+    client.loadUrl = async (url, selectorToBePresent) => {
       if (!options.chromeEnableAnimations) {
         debug('Disabling animations');
         await evaluateOnNewDocument(`(${disableAnimations})(window);`);
@@ -168,7 +175,19 @@ function createChromeTarget(
       await evaluateOnNewDocument(`(${disableInputCaret})(window);`);
 
       debug(`Navigating to ${url}`);
-      await Promise.all([Page.navigate({ url }), awaitRequestsFinished()]);
+      startObservingRequests();
+      await Page.navigate({ url });
+      await Page.loadEventFired();
+
+      if (selectorToBePresent) {
+        debug(`Awaiting selector "${selectorToBePresent}"`);
+        await executeFunctionWithWindow(
+          awaitSelectorPresent,
+          selectorToBePresent
+        );
+      }
+
+      await awaitRequestsFinished();
 
       debug('Awaiting runtime setup');
       await executeFunctionWithWindow(awaitLokiReady);
@@ -323,7 +342,7 @@ function createChromeTarget(
     const tab = await launchNewTab(tabOptions);
     let screenshot;
     try {
-      await withTimeout(options.chromeLoadTimeout)(tab.loadUrl(url));
+      await withTimeout(options.chromeLoadTimeout)(tab.loadUrl(url, selector));
       screenshot = await tab.captureScreenshot(selector);
     } catch (err) {
       if (err instanceof TimeoutError) {