From 2cc0f625080172f5d55f947eff1cdaa48cc6be98 Mon Sep 17 00:00:00 2001
From: Yann Braga <yannbf@gmail.com>
Date: Sat, 18 Dec 2021 12:09:26 +0100
Subject: [PATCH] sanitize passed urls and wait for storybook to load

---
 playwright/custom-environment.js | 75 +++++++++++++++++++++++++++-----
 1 file changed, 64 insertions(+), 11 deletions(-)

diff --git a/playwright/custom-environment.js b/playwright/custom-environment.js
index 05fb3edc..7ca284bb 100644
--- a/playwright/custom-environment.js
+++ b/playwright/custom-environment.js
@@ -1,30 +1,49 @@
 require('regenerator-runtime/runtime');
 const PlaywrightEnvironment = require('jest-playwright-preset/lib/PlaywrightEnvironment').default;
 
+const sanitizeURL = (url) => {
+  let finalURL = url
+  // prepend URL protocol if not there
+  if (finalURL.indexOf("http://") === -1 && finalURL.indexOf("https://") === -1) {
+    finalURL = 'http://' + finalURL;
+  }
+
+  // add forward slash at the end if not there
+  if (finalURL.slice(-1) !== '/') {
+    finalURL = finalURL + '/';
+  }
+
+  // remove iframe.html if present
+  return finalURL.replace(/iframe.html\s*$/, "");
+}
+
 class CustomEnvironment extends PlaywrightEnvironment {
   async setup() {
     await super.setup();
     const page = this.global.page;
     const start = new Date();
     const port = process.env.STORYBOOK_PORT || '6006';
-    const targetURL = process.env.TARGET_URL || `http://localhost:${port}`
-    const referenceURL = process.env.REFERENCE_URL
-    
-    if('TARGET_URL' in process.env && !process.env.TARGET_URL) {
-      console.log(`Received TARGET_URL but with a falsy value: ${
-        process.env.TARGET_URL
-      }, will fallback to ${targetURL} instead.`)
+    const targetURL = sanitizeURL(process.env.TARGET_URL || `http://localhost:${port}`);
+
+    const referenceURL = process.env.REFERENCE_URL && sanitizeURL(process.env.REFERENCE_URL);
+
+    if ('TARGET_URL' in process.env && !process.env.TARGET_URL) {
+      console.log(`Received TARGET_URL but with a falsy value: ${process.env.TARGET_URL
+        }, will fallback to ${targetURL} instead.`)
     }
 
-    await page.goto(`${targetURL}/iframe.html`, { waitUntil: 'load' }).catch((err) => {
-      if(err.message?.includes('ERR_CONNECTION_REFUSED')) {
+    await page.goto(`${targetURL}iframe.html`, { waitUntil: 'load' }).catch((err) => {
+      if (err.message?.includes('ERR_CONNECTION_REFUSED')) {
         const errorMessage = `Could not access the Storybook instance at ${targetURL}. Are you sure it's running?\n\n${err.message}`;
         throw new Error(errorMessage)
       }
 
       throw err;
     }); // FIXME: configure
-    console.log(`page loaded in ${new Date() - start}ms.`);
+    console.log(`page loaded in ${new Date() - start}ms.`, targetURL);
+
+    // if we ever want to log something from the browser to node
+    await page.exposeBinding('logToPage', (_, message) => console.log(message));
 
     await page.addScriptTag({
       content: `
@@ -43,13 +62,47 @@ class CustomEnvironment extends PlaywrightEnvironment {
           throw new StorybookTestRunnerError(storyId, errorMessage);
         }
 
+        async function __waitForElement(selector) {
+          return new Promise((resolve, reject) => {
+
+            const timeout = setTimeout(() => {
+              reject();
+            }, 10000);
+
+            if (document.querySelector(selector)) {
+              clearTimeout(timeout);
+              return resolve(document.querySelector(selector));
+            }
+
+            const observer = new MutationObserver(mutations => {
+              if (document.querySelector(selector)) {
+                clearTimeout(timeout);
+                resolve(document.querySelector(selector));
+                observer.disconnect();
+              }
+            });
+
+            observer.observe(document.body, {
+              childList: true,
+              subtree: true
+            });
+          });
+        }
+
         async function __test(storyId, hasPlayFn) {
+          try {
+            await __waitForElement('#root');
+          } catch(err) {
+            const message = \`Timed out waiting for Storybook to load after 10 seconds. Are you sure the Storybook is running correctly in that URL?\n\n\nHTML: \${document.body.innerHTML}\`;
+            throw new StorybookTestRunnerError(storyId, hasPlayFn, message);
+          }
+
           const channel = window.__STORYBOOK_ADDONS_CHANNEL__;
           if(!channel) {
             throw new StorybookTestRunnerError(
               storyId,
               hasPlayFn,
-              'The test runner could not access the story. Are you sure the Storybook is running correctly in that URL?'
+              'The test runner could not access the Storybook channel. Are you sure the Storybook is running correctly in that URL?'
             );
           }