Ability for custom storyshots testFunctions to utilise "snapshot per story file"

addons/storyshots/README.md

@@ -152,3 +152,29 @@ Like `snapshotWithOptions`, but generate a separate snapshot file for each stori
 ### `shallowSnapshot`
 
 Take a snapshot of a shallow-rendered version of the component.
+
+### `getSnapshotFileName`
+
+Utility function used in `multiSnapshotWithOptions`. This is made available for users who implement custom test functions that also want to take advantage of multi-file storyshots.
+
+###### Example:
+
+Let's say we wanted to create a test function for shallow && multi-file snapshots:
+
+```js
+import initStoryshots, { getSnapshotFileName } from '@storybook/addon-storyshots';
+import { shallow } from 'enzyme';
+import toJson from 'enzyme-to-json';
+
+initStoryshots({
+  test: ({ story, context }) => {
+    const snapshotFileName = getSnapshotFileName(context);
+    const storyElement = story.render(context);
+    const shallowTree = shallow(storyElement);
+
+    if (snapshotFileName) {
+      expect(toJson(shallowTree)).toMatchSpecificSnapshot(snapshotFileName);
+    }
+  }
+});
+```

addons/storyshots/src/index.js

@@ -8,7 +8,7 @@ import addons from '@storybook/addons';
 import runWithRequireContext from './require_context';
 import createChannel from './storybook-channel-mock';
 import { snapshot } from './test-bodies';
-import { getPossibleStoriesFiles } from './utils';
+import { getPossibleStoriesFiles, getSnapshotFileName } from './utils';
 
 export {
   snapshot,
@@ -18,6 +18,8 @@ export {
   renderOnly,
 } from './test-bodies';
 
+export { getSnapshotFileName };
+
 let storybook;
 let configPath;
 global.STORYBOOK_REACT_CLASSES = global.STORYBOOK_REACT_CLASSES || {};

addons/storyshots/src/test-bodies.js

@@ -1,23 +1,13 @@
 import renderer from 'react-test-renderer';
 import shallow from 'react-test-renderer/shallow';
 import 'jest-specific-snapshot';
-import { getStoryshotFile } from './utils';
+import { getSnapshotFileName } from './utils';
 
 function getRenderedTree(story, context, options) {
   const storyElement = story.render(context);
   return renderer.create(storyElement, options).toJSON();
 }
 
-function getSnapshotFileName(context) {
-  const fileName = context.fileName;
-
-  if (!fileName) {
-    return null;
-  }
-
-  return getStoryshotFile(fileName);
-}
-
 export const snapshotWithOptions = options => ({ story, context }) => {
   const tree = getRenderedTree(story, context, options);
   expect(tree).toMatchSnapshot();

addons/storyshots/src/utils.js

@@ -1,6 +1,6 @@
 import path from 'path';
 
-export function getStoryshotFile(fileName) {
+function getStoryshotFile(fileName) {
   const { dir, name } = path.parse(fileName);
   return path.format({ dir: path.join(dir, '__snapshots__'), name, ext: '.storyshot' });
 }
@@ -13,3 +13,13 @@ export function getPossibleStoriesFiles(storyshotFile) {
     path.format({ dir: path.dirname(dir), name, ext: '.jsx' }),
   ];
 }
+
+export function getSnapshotFileName(context) {
+  const fileName = context.fileName;
+
+  if (!fileName) {
+    return null;
+  }
+
+  return getStoryshotFile(fileName);
+}