From 6c971f256f01c7bfdcc9e04dd5899aae15d4aa30 Mon Sep 17 00:00:00 2001
From: Tom Coleman <tom@thesnail.org>
Date: Mon, 15 May 2017 15:38:43 +1000
Subject: [PATCH] Allow passing a `test` function to storyshots

See https://github.com/storybooks/storybook/issues/1034
---
 examples/test-cra/src/storyshots.test.js |  5 ++++-
 packages/storyshots/src/index.js         | 10 ++++++----
 packages/storyshots/src/test-bodies.js   | 13 +++++++++++++
 3 files changed, 23 insertions(+), 5 deletions(-)
 create mode 100644 packages/storyshots/src/test-bodies.js

diff --git a/examples/test-cra/src/storyshots.test.js b/examples/test-cra/src/storyshots.test.js
index 63f7047152ab..04ffd4761466 100644
--- a/examples/test-cra/src/storyshots.test.js
+++ b/examples/test-cra/src/storyshots.test.js
@@ -1,2 +1,5 @@
-import initStoryshots from 'storyshots';
+import initStoryshots, { renderOnly } from 'storyshots';
+
 initStoryshots();
+
+// initStoryshots({ test: renderOnly });
diff --git a/packages/storyshots/src/index.js b/packages/storyshots/src/index.js
index 5dc4d531ff10..e5f33130edcc 100644
--- a/packages/storyshots/src/index.js
+++ b/packages/storyshots/src/index.js
@@ -1,11 +1,13 @@
-import renderer from 'react-test-renderer';
 import path from 'path';
 import readPkgUp from 'read-pkg-up';
 import addons from '@kadira/storybook-addons';
 import runWithRequireContext from './require_context';
 import createChannel from './storybook-channel-mock';
+import { snapshot } from './test-bodies';
 const { describe, it, expect } = global;
 
+export { snapshot, renderOnly } from './test-bodies';
+
 let storybook;
 let configPath;
 
@@ -60,6 +62,8 @@ export default function testStorySnapshots(options = {}) {
   // Added not to break existing storyshots configs (can be removed in a future major release)
   options.storyNameRegex = options.storyNameRegex || options.storyRegex;
 
+  options.test = options.test || snapshot;
+
   for (const group of stories) {
     if (options.storyKindRegex && !group.kind.match(options.storyKindRegex)) {
       continue;
@@ -74,9 +78,7 @@ export default function testStorySnapshots(options = {}) {
 
           it(story.name, () => {
             const context = { kind: group.kind, story: story.name };
-            const renderedStory = story.render(context);
-            const tree = renderer.create(renderedStory).toJSON();
-            expect(tree).toMatchSnapshot();
+            options.test({ story, context });
           });
         }
       });
diff --git a/packages/storyshots/src/test-bodies.js b/packages/storyshots/src/test-bodies.js
new file mode 100644
index 000000000000..0da64358a671
--- /dev/null
+++ b/packages/storyshots/src/test-bodies.js
@@ -0,0 +1,13 @@
+import renderer from 'react-test-renderer';
+import shallow from 'react-test-renderer/shallow';
+
+export function snapshot({ story, context }) {
+  const storyElement = story.render(context);
+  const tree = renderer.create(storyElement).toJSON();
+  expect(tree).toMatchSnapshot();
+}
+
+export function renderOnly({ story, context }) {
+  const storyElement = story.render(context);
+  const tree = renderer.create(storyElement);
+}