Storyshots: Fix async issue with obtaining custom Puppeteer instance (#…

…5069#)

addons/storyshots/storyshots-puppeteer/README.md

@@ -148,19 +148,20 @@ initStoryshots({
 
 ### Specifying a custom puppeteer `browser` instance
 
-You might use `customBrowser` to specify a custom instance of a puppeteer `browser` object. This will prevent `storyshots-puppeteer` from creating its own `browser`. It will create and close pages within the `browser`, and it is your responsibility to manage the lifecycle of the `browser` itself.
+You might use the async `getCustomBrowser` function to obtain a custom instance of a puppeteer `browser` object. This will prevent `storyshots-puppeteer` from creating its own `browser`. It will create and close pages within the `browser`, and it is your responsibility to manage the lifecycle of the `browser` itself.
 
 ```js
 import initStoryshots from '@storybook/addon-storyshots';
 import { imageSnapshot } from '@storybook/addon-storyshots-puppeteer';
 import puppeteer from 'puppeteer';
 
 (async function() {
-    const customBrowser = await puppeteer.connect('ws://yourUrl');
-
     initStoryshots({
       suite: 'Image storyshots',
-      test: imageSnapshot({ storybookUrl: 'http://localhost:6006', customBrowser }),
+      test: imageSnapshot({ 
+        storybookUrl: 'http://localhost:6006', 
+        getCustomBrowser: async () => puppeteer.connect('ws://yourUrl')
+      }),
     });
 })();
 ```

addons/storyshots/storyshots-puppeteer/src/index.js

@@ -19,7 +19,7 @@ const defaultConfig = {
   beforeScreenshot: noop,
   getGotoOptions: noop,
   customizePage: asyncNoop,
-  customBrowser: undefined,
+  getCustomBrowser: asyncNoop,
 };
 
 export const imageSnapshot = (customConfig = {}) => {
@@ -31,16 +31,12 @@ export const imageSnapshot = (customConfig = {}) => {
     beforeScreenshot,
     getGotoOptions,
     customizePage,
-    customBrowser,
+    getCustomBrowser,
   } = { ...defaultConfig, ...customConfig };
 
   let browser; // holds ref to browser. (ie. Chrome)
   let page; // Hold ref to the page to screenshot.
 
-  if (customBrowser) {
-    browser = customBrowser;
-  }
-
   const testFn = async ({ context }) => {
     const { kind, framework, story } = context;
     if (framework === 'rn') {
@@ -81,15 +77,17 @@ export const imageSnapshot = (customConfig = {}) => {
   };
 
   testFn.afterAll = () => {
-    if (customBrowser) {
+    if (getCustomBrowser && page) {
       return page.close();
     }
 
     return browser.close();
   };
 
   testFn.beforeAll = async () => {
-    if (!browser) {
+    if (getCustomBrowser) {
+      browser = await getCustomBrowser();
+    } else {
       // add some options "no-sandbox" to make it work properly on some Linux systems as proposed here: https://github.com/Googlechrome/puppeteer/issues/290#issuecomment-322851507
       browser = await puppeteer.launch({
         args: ['--no-sandbox ', '--disable-setuid-sandbox'],