From 2512cc2786c5bb546ca58b93d75b9fd2754f65db Mon Sep 17 00:00:00 2001
From: Michael Shilman <michael@lab80.co>
Date: Mon, 7 Oct 2019 16:07:22 +0800
Subject: [PATCH 1/5] MDX: Support imported story functions

---
 .../mdx-compiler-plugin.test.js.snap          | 91 ++++++++++++++++++-
 .../__testfixtures__/story-function-var.mdx   | 11 +++
 addons/docs/src/mdx/mdx-compiler-plugin.js    | 10 +-
 .../docs/src/mdx/mdx-compiler-plugin.test.js  |  1 +
 4 files changed, 106 insertions(+), 7 deletions(-)
 create mode 100644 addons/docs/src/mdx/__testfixtures__/story-function-var.mdx

diff --git a/addons/docs/src/mdx/__snapshots__/mdx-compiler-plugin.test.js.snap b/addons/docs/src/mdx/__snapshots__/mdx-compiler-plugin.test.js.snap
index 519c6c08745b..2f1537fe85b1 100644
--- a/addons/docs/src/mdx/__snapshots__/mdx-compiler-plugin.test.js.snap
+++ b/addons/docs/src/mdx/__snapshots__/mdx-compiler-plugin.test.js.snap
@@ -51,6 +51,8 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
+const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
+
 export const one = () => <Button>One</Button>;
 one.story = {};
 one.story.name = 'one';
@@ -121,6 +123,8 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
+const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
+
 export const __page = () => {
   throw new Error('Docs-only story');
 };
@@ -182,6 +186,8 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
+const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
+
 export const one = () => <Button>One</Button>;
 one.story = {};
 one.story.name = 'one';
@@ -256,6 +262,8 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
+const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
+
 export const componentNotes = () => <Button>Component notes</Button>;
 componentNotes.story = {};
 componentNotes.story.name = 'component notes';
@@ -344,6 +352,8 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
+const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
+
 export const helloButton = () => <Button>Hello button</Button>;
 helloButton.story = {};
 helloButton.story.name = 'hello button';
@@ -405,6 +415,8 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
+const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
+
 const componentMeta = { includeStories: [] };
 
 const mdxStoryNameToId = {};
@@ -453,7 +465,9 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
-export const text = () => 'Plain text';
+const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
+
+export const text = makeStoryFn('Plain text');
 text.story = {};
 text.story.name = 'text';
 text.story.parameters = { mdxSource: \\"'Plain text'\\" };
@@ -513,6 +527,8 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
+const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
+
 export const one = () => <Button>One</Button>;
 one.story = {};
 one.story.name = 'one';
@@ -581,12 +597,14 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
-export const functionStory = () => {
+const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
+
+export const functionStory = makeStoryFn(() => {
   const btn = document.createElement('button');
   btn.innerHTML = 'Hello Button';
   btn.addEventListener('click', action('Click'));
   return btn;
-};
+});
 functionStory.story = {};
 functionStory.story.name = 'function';
 functionStory.story.parameters = {
@@ -610,6 +628,65 @@ export default componentMeta;
 "
 `;
 
+exports[`docs-mdx-compiler-plugin story-function-var.mdx 1`] = `
+"/* @jsx mdx */
+import { DocsContainer } from '@storybook/addon-docs/blocks';
+
+import { Meta, Story } from '@storybook/addon-docs/blocks';
+export const basicFn = () => <Button mdxType=\\"Button\\" />;
+const makeShortcode = name =>
+  function MDXDefaultShortcode(props) {
+    console.warn(
+      'Component ' +
+        name +
+        ' was not imported, exported, or provided by MDXProvider as global scope'
+    );
+    return <div {...props} />;
+  };
+const Button = makeShortcode('Button');
+const layoutProps = {
+  basicFn,
+};
+const MDXLayout = 'wrapper';
+function MDXContent({ components, ...props }) {
+  return (
+    <MDXLayout {...layoutProps} {...props} components={components} mdxType=\\"MDXLayout\\">
+      <Meta title=\\"story-function-var\\" mdxType=\\"Meta\\" />
+
+      <h1>{\`Button\`}</h1>
+      <p>{\`I can define a story with the function defined in CSF:\`}</p>
+      <Story name=\\"basic\\" mdxType=\\"Story\\">
+        {basicFn}
+      </Story>
+    </MDXLayout>
+  );
+}
+
+MDXContent.isMDXComponent = true;
+
+const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
+
+export const basic = makeStoryFn(basicFn);
+basic.story = {};
+basic.story.name = 'basic';
+basic.story.parameters = { mdxSource: 'basicFn' };
+
+const componentMeta = { title: 'story-function-var', includeStories: ['basic'] };
+
+const mdxStoryNameToId = { basic: 'story-function-var--basic' };
+
+componentMeta.parameters = componentMeta.parameters || {};
+componentMeta.parameters.docs = {
+  container: ({ context, children }) => (
+    <DocsContainer context={{ ...context, mdxStoryNameToId }}>{children}</DocsContainer>
+  ),
+  page: MDXContent,
+};
+
+export default componentMeta;
+"
+`;
+
 exports[`docs-mdx-compiler-plugin story-object.mdx 1`] = `
 "/* @jsx mdx */
 import { DocsContainer } from '@storybook/addon-docs/blocks';
@@ -652,7 +729,9 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
-export const toStorybook = () => ({
+const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
+
+export const toStorybook = makeStoryFn({
   template: \`<storybook-welcome-component (showApp)=\\"showApp()\\"></storybook-welcome-component>\`,
   props: {
     showApp: linkTo('Button'),
@@ -713,6 +792,8 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
+const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
+
 const componentMeta = { includeStories: [] };
 
 const mdxStoryNameToId = {};
@@ -759,6 +840,8 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
+const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
+
 const componentMeta = { includeStories: [] };
 
 const mdxStoryNameToId = {};
diff --git a/addons/docs/src/mdx/__testfixtures__/story-function-var.mdx b/addons/docs/src/mdx/__testfixtures__/story-function-var.mdx
new file mode 100644
index 000000000000..7cfbae252f7b
--- /dev/null
+++ b/addons/docs/src/mdx/__testfixtures__/story-function-var.mdx
@@ -0,0 +1,11 @@
+import { Meta, Story } from '@storybook/addon-docs/blocks';
+
+<Meta title="story-function-var" />
+
+export const basicFn = () => <Button />;
+
+# Button
+
+I can define a story with the function defined in CSF:
+
+<Story name="basic">{basicFn}</Story>
diff --git a/addons/docs/src/mdx/mdx-compiler-plugin.js b/addons/docs/src/mdx/mdx-compiler-plugin.js
index 5414993eac8d..572526d53f01 100644
--- a/addons/docs/src/mdx/mdx-compiler-plugin.js
+++ b/addons/docs/src/mdx/mdx-compiler-plugin.js
@@ -52,6 +52,7 @@ function genStoryExport(ast, context) {
 
   let body = ast.children.find(n => n.type !== 'JSXText');
   let storyCode = null;
+  let isJsx = false;
   if (!body) {
     // plain text node
     const { code } = generate(ast.children[0], {});
@@ -60,18 +61,20 @@ function genStoryExport(ast, context) {
     if (body.type === 'JSXExpressionContainer') {
       // FIXME: handle fragments
       body = body.expression;
+    } else {
+      isJsx = true;
     }
     const { code } = generate(body, {});
     storyCode = code;
   }
-  if (storyCode.trim().startsWith('() =>')) {
-    statements.push(`export const ${storyKey} = ${storyCode}`);
-  } else {
+  if (isJsx) {
     statements.push(
       `export const ${storyKey} = () => (
         ${storyCode}
       );`
     );
+  } else {
+    statements.push(`export const ${storyKey} = makeStoryFn(${storyCode});`);
   }
   statements.push(`${storyKey}.story = {};`);
 
@@ -242,6 +245,7 @@ function extractExports(node, options) {
   const fullJsx = [
     'import { DocsContainer } from "@storybook/addon-docs/blocks";',
     defaultJsx,
+    `const makeStoryFn = (val) => (typeof val === 'function' ? val : () => val);`,
     ...storyExports,
     `const componentMeta = ${stringifyMeta(metaExport)};`,
     `const mdxStoryNameToId = ${JSON.stringify(mdxStoryNameToId)};`,
diff --git a/addons/docs/src/mdx/mdx-compiler-plugin.test.js b/addons/docs/src/mdx/mdx-compiler-plugin.test.js
index a88aace02d4c..81d9a4a779be 100644
--- a/addons/docs/src/mdx/mdx-compiler-plugin.test.js
+++ b/addons/docs/src/mdx/mdx-compiler-plugin.test.js
@@ -36,6 +36,7 @@ describe('docs-mdx-compiler-plugin', () => {
     'non-story-exports.mdx',
     'story-function.mdx',
     'docs-only.mdx',
+    'story-function-var.mdx',
   ];
   fixtures.forEach(fixtureFile => {
     it(fixtureFile, async () => {

From 0d9b0f973538f7727ed1cca134c378571fbc2d53 Mon Sep 17 00:00:00 2001
From: Michael Shilman <michael@lab80.co>
Date: Mon, 7 Oct 2019 16:07:58 +0800
Subject: [PATCH 2/5] MDX: Better recipe for MDX documentation of CSF + example

---
 addons/docs/docs/recipes.md                   | 56 ++++++++++---------
 .../addon-docs/csf-with-mdx-docs.stories.js   | 11 ++++
 .../addon-docs/csf-with-mdx-docs.stories.mdx  | 10 ++++
 3 files changed, 51 insertions(+), 26 deletions(-)
 create mode 100644 examples/official-storybook/stories/addon-docs/csf-with-mdx-docs.stories.js
 create mode 100644 examples/official-storybook/stories/addon-docs/csf-with-mdx-docs.stories.mdx

diff --git a/addons/docs/docs/recipes.md b/addons/docs/docs/recipes.md
index 9c18bb74a5a3..fd3ae384857b 100644
--- a/addons/docs/docs/recipes.md
+++ b/addons/docs/docs/recipes.md
@@ -31,46 +31,50 @@ The only limitation is that your exported titles (CSF: `default.title`, MDX `Met
 
 Perhaps you want to write your stories in CSF, but document them in MDX? Here's how to do that:
 
-**Button.mdx**
+**Button.stories.js**
+
+```js
+import React from 'react';
+import { Button } from './Button';
+
+export default {
+  title: 'Demo/Button',
+  component: Button,
+  includeStories: [], // or simply don't load this file at all
+};
+
+export const basic = () => <Button>Basic</Button>;
+basic.story = {
+  parameters: { foo: 'bar' },
+};
+```
+
+**Button.stories.mdx**
 
 ```md
-import { Story } from '@storybook/addon-docs/blocks';
-import { SomeComponent } from 'somewhere';
+import { Meta, Story } from '@storybook/addon-docs/blocks';
+import \* as stories from './Button.stories.js';
+import { SomeComponent } from 'path/to/SomeComponent';
+
+<Meta {...stories.default} />
 
 # Button
 
 I can embed a story (but not define one, since this file should not contain a `Meta`):
 
-<Story id="some--id" />
+<Story name="basic">{stories.basic}</Story>
 
 And of course I can also embed arbitrary markdown & JSX in this file.
 
 <SomeComponent prop1="val1" />
 ```
 
-**Button.stories.js**
-
-```js
-import { Button } from './Button';
-import mdx from './Button.mdx';
-
-export default {
-  title: 'Demo/Button',
-  parameters: {
-    docs: {
-      page: mdx,
-    },
-  },
-};
-
-export const basic = () => <Button>Basic</Button>;
-```
-
-Note that in contrast to other examples, the MDX file suffix is `.mdx` rather than `.stories.mdx`. This key difference means that the file will be loaded with the default MDX loader rather than Storybook's CSF loader, which has several implications:
+What's happening here:
 
-1. You don't need to provide a `Meta` declaration.
-2. You can refer to existing stories (i.e. `<Story id="...">`) but cannot define new stories (i.e. `<Story name="...">`).
-3. The documentation gets exported as the default export (MDX default) rather than as a parameter hanging off the default export (CSF).
+- Your stories are defined in CSF, but because of `includeStories: []`, they are not actually added to Storybook.
+- The MDX file is adding the stories to Storybook, and using the story function defined in CSF.
+- The MDX loader is using story metadata from CSF, such as name, decorators, parameters, but will give giving preference to anything defined in the MDX file.
+- The MDX file is using the Meta `default` defined in the CSF.
 
 ## Mixing storiesOf with CSF/MDX
 
diff --git a/examples/official-storybook/stories/addon-docs/csf-with-mdx-docs.stories.js b/examples/official-storybook/stories/addon-docs/csf-with-mdx-docs.stories.js
new file mode 100644
index 000000000000..655442485112
--- /dev/null
+++ b/examples/official-storybook/stories/addon-docs/csf-with-mdx-docs.stories.js
@@ -0,0 +1,11 @@
+import React from 'react';
+import { Button } from '@storybook/react/demo';
+
+export default {
+  title: 'Addons|Docs/csf-with-mdx-docs',
+  component: Button,
+  includeStories: [], // or simply don't load this file at all
+};
+
+// eslint-disable-next-line react/prop-types
+export const basic = ({ parameters }) => <Button>Basic</Button>;
diff --git a/examples/official-storybook/stories/addon-docs/csf-with-mdx-docs.stories.mdx b/examples/official-storybook/stories/addon-docs/csf-with-mdx-docs.stories.mdx
new file mode 100644
index 000000000000..16472c77619b
--- /dev/null
+++ b/examples/official-storybook/stories/addon-docs/csf-with-mdx-docs.stories.mdx
@@ -0,0 +1,10 @@
+import { Meta, Story } from '@storybook/addon-docs/blocks';
+import * as stories from './csf-with-mdx-docs.stories';
+
+<Meta title="Addons|Docs/csf-with-mdx-docs" />
+
+# Button
+
+I can define a story with the function defined in CSF!
+
+<Story name="basic">{stories.basic}</Story>

From 6fa7e2662a5747047a40265ef4aeb5632486ead0 Mon Sep 17 00:00:00 2001
From: Michael Shilman <shilman@users.noreply.github.com>
Date: Mon, 7 Oct 2019 18:28:31 +0800
Subject: [PATCH 3/5] Update addons/docs/docs/recipes.md

Co-Authored-By: Tom Coleman <tom@thesnail.org>
---
 addons/docs/docs/recipes.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/addons/docs/docs/recipes.md b/addons/docs/docs/recipes.md
index fd3ae384857b..064514fe48fa 100644
--- a/addons/docs/docs/recipes.md
+++ b/addons/docs/docs/recipes.md
@@ -53,7 +53,7 @@ basic.story = {
 
 ```md
 import { Meta, Story } from '@storybook/addon-docs/blocks';
-import \* as stories from './Button.stories.js';
+import * as stories from './Button.stories.js';
 import { SomeComponent } from 'path/to/SomeComponent';
 
 <Meta {...stories.default} />

From 887775729979bdcd926815042f498fb6c1c5c412 Mon Sep 17 00:00:00 2001
From: Michael Shilman <michael@lab80.co>
Date: Mon, 7 Oct 2019 18:35:45 +0800
Subject: [PATCH 4/5] Fix typo

---
 addons/docs/docs/recipes.md                                     | 2 +-
 .../stories/addon-docs/csf-with-mdx-docs.stories.mdx            | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/addons/docs/docs/recipes.md b/addons/docs/docs/recipes.md
index fd3ae384857b..ca0f8a021f5d 100644
--- a/addons/docs/docs/recipes.md
+++ b/addons/docs/docs/recipes.md
@@ -60,7 +60,7 @@ import { SomeComponent } from 'path/to/SomeComponent';
 
 # Button
 
-I can embed a story (but not define one, since this file should not contain a `Meta`):
+I can define a story with the function imported from CSF:
 
 <Story name="basic">{stories.basic}</Story>
 
diff --git a/examples/official-storybook/stories/addon-docs/csf-with-mdx-docs.stories.mdx b/examples/official-storybook/stories/addon-docs/csf-with-mdx-docs.stories.mdx
index 16472c77619b..485a6dbd6b49 100644
--- a/examples/official-storybook/stories/addon-docs/csf-with-mdx-docs.stories.mdx
+++ b/examples/official-storybook/stories/addon-docs/csf-with-mdx-docs.stories.mdx
@@ -5,6 +5,6 @@ import * as stories from './csf-with-mdx-docs.stories';
 
 # Button
 
-I can define a story with the function defined in CSF!
+I can define a story with the function imported from CSF:
 
 <Story name="basic">{stories.basic}</Story>

From 717665daf13eebe743c4d03487762ce03b99f93d Mon Sep 17 00:00:00 2001
From: Michael Shilman <michael@lab80.co>
Date: Mon, 7 Oct 2019 18:37:57 +0800
Subject: [PATCH 5/5] MDX: Move makeStoryFn into library

---
 addons/docs/src/blocks/index.ts               |  3 ++
 .../mdx-compiler-plugin.test.js.snap          | 52 +++++--------------
 addons/docs/src/mdx/mdx-compiler-plugin.js    |  3 +-
 3 files changed, 17 insertions(+), 41 deletions(-)

diff --git a/addons/docs/src/blocks/index.ts b/addons/docs/src/blocks/index.ts
index 96a505258171..4e68cf50276c 100644
--- a/addons/docs/src/blocks/index.ts
+++ b/addons/docs/src/blocks/index.ts
@@ -11,3 +11,6 @@ export * from './Props';
 export * from './Source';
 export * from './Story';
 export * from './Wrapper';
+
+// helper function for MDX
+export const makeStoryFn = (val: any) => (typeof val === 'function' ? val : () => val);
diff --git a/addons/docs/src/mdx/__snapshots__/mdx-compiler-plugin.test.js.snap b/addons/docs/src/mdx/__snapshots__/mdx-compiler-plugin.test.js.snap
index 2f1537fe85b1..2456a29c4ee3 100644
--- a/addons/docs/src/mdx/__snapshots__/mdx-compiler-plugin.test.js.snap
+++ b/addons/docs/src/mdx/__snapshots__/mdx-compiler-plugin.test.js.snap
@@ -2,7 +2,7 @@
 
 exports[`docs-mdx-compiler-plugin decorators.mdx 1`] = `
 "/* @jsx mdx */
-import { DocsContainer } from '@storybook/addon-docs/blocks';
+import { DocsContainer, makeStoryFn } from '@storybook/addon-docs/blocks';
 
 import { Button } from '@storybook/react/demo';
 import { Story, Meta } from '@storybook/addon-docs/blocks';
@@ -51,8 +51,6 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
-const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
-
 export const one = () => <Button>One</Button>;
 one.story = {};
 one.story.name = 'one';
@@ -91,7 +89,7 @@ export default componentMeta;
 
 exports[`docs-mdx-compiler-plugin docs-only.mdx 1`] = `
 "/* @jsx mdx */
-import { DocsContainer } from '@storybook/addon-docs/blocks';
+import { DocsContainer, makeStoryFn } from '@storybook/addon-docs/blocks';
 
 import { Meta } from '@storybook/addon-docs/blocks';
 
@@ -123,8 +121,6 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
-const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
-
 export const __page = () => {
   throw new Error('Docs-only story');
 };
@@ -149,7 +145,7 @@ export default componentMeta;
 
 exports[`docs-mdx-compiler-plugin non-story-exports.mdx 1`] = `
 "/* @jsx mdx */
-import { DocsContainer } from '@storybook/addon-docs/blocks';
+import { DocsContainer, makeStoryFn } from '@storybook/addon-docs/blocks';
 
 import { Button } from '@storybook/react/demo';
 import { Story, Meta } from '@storybook/addon-docs/blocks';
@@ -186,8 +182,6 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
-const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
-
 export const one = () => <Button>One</Button>;
 one.story = {};
 one.story.name = 'one';
@@ -216,7 +210,7 @@ export default componentMeta;
 
 exports[`docs-mdx-compiler-plugin parameters.mdx 1`] = `
 "/* @jsx mdx */
-import { DocsContainer } from '@storybook/addon-docs/blocks';
+import { DocsContainer, makeStoryFn } from '@storybook/addon-docs/blocks';
 
 import { Button } from '@storybook/react/demo';
 import { Story, Meta } from '@storybook/addon-docs/blocks';
@@ -262,8 +256,6 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
-const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
-
 export const componentNotes = () => <Button>Component notes</Button>;
 componentNotes.story = {};
 componentNotes.story.name = 'component notes';
@@ -306,7 +298,7 @@ export default componentMeta;
 
 exports[`docs-mdx-compiler-plugin previews.mdx 1`] = `
 "/* @jsx mdx */
-import { DocsContainer } from '@storybook/addon-docs/blocks';
+import { DocsContainer, makeStoryFn } from '@storybook/addon-docs/blocks';
 
 import { Button } from '@storybook/react/demo';
 import { Preview, Story, Meta } from '@storybook/addon-docs/blocks';
@@ -352,8 +344,6 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
-const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
-
 export const helloButton = () => <Button>Hello button</Button>;
 helloButton.story = {};
 helloButton.story.name = 'hello button';
@@ -388,7 +378,7 @@ export default componentMeta;
 
 exports[`docs-mdx-compiler-plugin story-current.mdx 1`] = `
 "/* @jsx mdx */
-import { DocsContainer } from '@storybook/addon-docs/blocks';
+import { DocsContainer, makeStoryFn } from '@storybook/addon-docs/blocks';
 
 import { Story } from '@storybook/addon-docs/blocks';
 
@@ -415,8 +405,6 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
-const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
-
 const componentMeta = { includeStories: [] };
 
 const mdxStoryNameToId = {};
@@ -435,7 +423,7 @@ export default componentMeta;
 
 exports[`docs-mdx-compiler-plugin story-def-text-only.mdx 1`] = `
 "/* @jsx mdx */
-import { DocsContainer } from '@storybook/addon-docs/blocks';
+import { DocsContainer, makeStoryFn } from '@storybook/addon-docs/blocks';
 
 import { Story, Meta } from '@storybook/addon-docs/blocks';
 
@@ -465,8 +453,6 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
-const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
-
 export const text = makeStoryFn('Plain text');
 text.story = {};
 text.story.name = 'text';
@@ -490,7 +476,7 @@ export default componentMeta;
 
 exports[`docs-mdx-compiler-plugin story-definitions.mdx 1`] = `
 "/* @jsx mdx */
-import { DocsContainer } from '@storybook/addon-docs/blocks';
+import { DocsContainer, makeStoryFn } from '@storybook/addon-docs/blocks';
 
 import { Button } from '@storybook/react/demo';
 import { Story, Meta } from '@storybook/addon-docs/blocks';
@@ -527,8 +513,6 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
-const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
-
 export const one = () => <Button>One</Button>;
 one.story = {};
 one.story.name = 'one';
@@ -566,7 +550,7 @@ export default componentMeta;
 
 exports[`docs-mdx-compiler-plugin story-function.mdx 1`] = `
 "/* @jsx mdx */
-import { DocsContainer } from '@storybook/addon-docs/blocks';
+import { DocsContainer, makeStoryFn } from '@storybook/addon-docs/blocks';
 
 const makeShortcode = name =>
   function MDXDefaultShortcode(props) {
@@ -597,8 +581,6 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
-const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
-
 export const functionStory = makeStoryFn(() => {
   const btn = document.createElement('button');
   btn.innerHTML = 'Hello Button';
@@ -630,7 +612,7 @@ export default componentMeta;
 
 exports[`docs-mdx-compiler-plugin story-function-var.mdx 1`] = `
 "/* @jsx mdx */
-import { DocsContainer } from '@storybook/addon-docs/blocks';
+import { DocsContainer, makeStoryFn } from '@storybook/addon-docs/blocks';
 
 import { Meta, Story } from '@storybook/addon-docs/blocks';
 export const basicFn = () => <Button mdxType=\\"Button\\" />;
@@ -664,8 +646,6 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
-const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
-
 export const basic = makeStoryFn(basicFn);
 basic.story = {};
 basic.story.name = 'basic';
@@ -689,7 +669,7 @@ export default componentMeta;
 
 exports[`docs-mdx-compiler-plugin story-object.mdx 1`] = `
 "/* @jsx mdx */
-import { DocsContainer } from '@storybook/addon-docs/blocks';
+import { DocsContainer, makeStoryFn } from '@storybook/addon-docs/blocks';
 
 import { Story, Meta } from '@storybook/addon-docs/blocks';
 import { Welcome, Button } from '@storybook/angular/demo';
@@ -729,8 +709,6 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
-const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
-
 export const toStorybook = makeStoryFn({
   template: \`<storybook-welcome-component (showApp)=\\"showApp()\\"></storybook-welcome-component>\`,
   props: {
@@ -765,7 +743,7 @@ export default componentMeta;
 
 exports[`docs-mdx-compiler-plugin story-references.mdx 1`] = `
 "/* @jsx mdx */
-import { DocsContainer } from '@storybook/addon-docs/blocks';
+import { DocsContainer, makeStoryFn } from '@storybook/addon-docs/blocks';
 
 import { Story } from '@storybook/addon-docs/blocks';
 
@@ -792,8 +770,6 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
-const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
-
 const componentMeta = { includeStories: [] };
 
 const mdxStoryNameToId = {};
@@ -812,7 +788,7 @@ export default componentMeta;
 
 exports[`docs-mdx-compiler-plugin vanilla.mdx 1`] = `
 "/* @jsx mdx */
-import { DocsContainer } from '@storybook/addon-docs/blocks';
+import { DocsContainer, makeStoryFn } from '@storybook/addon-docs/blocks';
 
 import { Button } from '@storybook/react/demo';
 
@@ -840,8 +816,6 @@ function MDXContent({ components, ...props }) {
 
 MDXContent.isMDXComponent = true;
 
-const makeStoryFn = val => (typeof val === 'function' ? val : () => val);
-
 const componentMeta = { includeStories: [] };
 
 const mdxStoryNameToId = {};
diff --git a/addons/docs/src/mdx/mdx-compiler-plugin.js b/addons/docs/src/mdx/mdx-compiler-plugin.js
index 572526d53f01..84aa135fdfca 100644
--- a/addons/docs/src/mdx/mdx-compiler-plugin.js
+++ b/addons/docs/src/mdx/mdx-compiler-plugin.js
@@ -243,9 +243,8 @@ function extractExports(node, options) {
   );
 
   const fullJsx = [
-    'import { DocsContainer } from "@storybook/addon-docs/blocks";',
+    'import { DocsContainer, makeStoryFn } from "@storybook/addon-docs/blocks";',
     defaultJsx,
-    `const makeStoryFn = (val) => (typeof val === 'function' ? val : () => val);`,
     ...storyExports,
     `const componentMeta = ${stringifyMeta(metaExport)};`,
     `const mdxStoryNameToId = ${JSON.stringify(mdxStoryNameToId)};`,