Allow promise to be returned from 'add()'

addons/storyshots/src/index.js

@@ -86,7 +86,7 @@ export default function testStorySnapshots(options = {}) {
 
           it(story.name, () => {
             const context = { kind: group.kind, story: story.name };
-            options.test({ story, context });
+            return options.test({ story, context });
           });
         }
       });

addons/storyshots/src/test-bodies.js

@@ -1,11 +1,11 @@
 import renderer from 'react-test-renderer';
 import shallow from 'react-test-renderer/shallow';
 
-export const snapshotWithOptions = options => ({ story, context }) => {
-  const storyElement = story.render(context);
-  const tree = renderer.create(storyElement, options).toJSON();
-  expect(tree).toMatchSnapshot();
-};
+export const snapshotWithOptions = options => ({ story, context }) =>
+  Promise.resolve(story).then(storyVal => storyVal.render(context)).then(storyElement => {
+    const tree = renderer.create(storyElement, options).toJSON();
+    expect(tree).toMatchSnapshot();
+  });
 
 export const snapshot = snapshotWithOptions({});
 

app/react/src/client/preview/render.js

@@ -71,32 +71,32 @@ export function renderMain(data, storyStore) {
     story: selectedStory,
   };
 
-  const element = story(context);
-
-  if (!element) {
-    const error = {
-      title: `Expecting a React element from the story: "${selectedStory}" of "${selectedKind}".`,
-      description: stripIndents`
-        Did you forget to return the React element from the story?
-        Use "() => (<MyComp/>)" or "() => { return <MyComp/>; }" when defining the story.
-      `,
-    };
-    return renderError(error);
-  }
-
-  if (!isReactRenderable(element)) {
-    const error = {
-      title: `Expecting a valid React element from the story: "${selectedStory}" of "${selectedKind}".`,
-      description: stripIndents`
-         Seems like you are not returning a correct React element from the story.
-         Could you double check that?
-       `,
-    };
-    return renderError(error);
-  }
-
-  ReactDOM.render(element, rootEl);
-  return null;
+  return Promise.resolve(story(context)).then(element => {
+    if (!element) {
+      const error = {
+        title: `Expecting a React element from the story: "${selectedStory}" of "${selectedKind}".`,
+        description: stripIndents`
+          Did you forget to return the React element from the story?
+          Use "() => (<MyComp/>)" or "() => { return <MyComp/>; }" when defining the story.
+        `,
+      };
+      return renderError(error);
+    }
+
+    if (!isReactRenderable(element)) {
+      const error = {
+        title: `Expecting a valid React element from the story: "${selectedStory}" of "${selectedKind}".`,
+        description: stripIndents`
+           Seems like you are not returning a correct React element from the story.
+           Could you double check that?
+         `,
+      };
+      return renderError(error);
+    }
+
+    ReactDOM.render(element, rootEl);
+    return null;
+  });
 }
 
 export default function renderPreview({ reduxStore, storyStore }) {