From f4ace6c009615a1459910d71cce5a9ccae6deb10 Mon Sep 17 00:00:00 2001
From: Nikki Reece <nicole.reece@monsanto.com>
Date: Mon, 31 Aug 2020 10:59:54 -0500
Subject: [PATCH 1/3] send context to transformSource

---
 addons/docs/docs/recipes.md                   |  2 +-
 addons/docs/src/blocks/enhanceSource.test.ts  | 35 +++++++++++++++++--
 addons/docs/src/blocks/enhanceSource.ts       |  2 +-
 .../__snapshots__/button.stories.storyshot    |  6 ++++
 .../stories/button.stories.js                 | 35 +++++++++++++++----
 lib/api/src/version.ts                        |  2 +-
 lib/cli/versions.json                         |  2 +-
 7 files changed, 72 insertions(+), 12 deletions(-)

diff --git a/addons/docs/docs/recipes.md b/addons/docs/docs/recipes.md
index e08fe60c0f0e..9eadd81ba218 100644
--- a/addons/docs/docs/recipes.md
+++ b/addons/docs/docs/recipes.md
@@ -270,7 +270,7 @@ Alternatively, you can provide a function in the `docs.transformSource` paramete
 const SOURCE_REGEX = /^\(\) => `(.*)`$/;
 export const parameters = {
   docs: {
-    transformSource: (src, storyId) => {
+    transformSource: (src, storyId, storyContext) => {
       const match = SOURCE_REGEX.exec(src);
       return match ? match[1] : src;
     },
diff --git a/addons/docs/src/blocks/enhanceSource.test.ts b/addons/docs/src/blocks/enhanceSource.test.ts
index 4f977884858a..695e66abad1d 100644
--- a/addons/docs/src/blocks/enhanceSource.test.ts
+++ b/addons/docs/src/blocks/enhanceSource.test.ts
@@ -11,8 +11,6 @@ const emptyContext: StoryContext = {
   parameters: {},
 };
 
-const transformSource = (src?: string) => (src ? `formatted: ${src}` : 'no src');
-
 describe('addon-docs enhanceSource', () => {
   describe('no source loaded', () => {
     const baseContext = emptyContext;
@@ -20,6 +18,7 @@ describe('addon-docs enhanceSource', () => {
       expect(enhanceSource(baseContext)).toBeNull();
     });
     it('transformSource', () => {
+      const transformSource = (src?: string) => (src ? `formatted: ${src}` : 'no src');
       const parameters = { ...baseContext.parameters, docs: { transformSource } };
       expect(enhanceSource({ ...baseContext, parameters })).toBeNull();
     });
@@ -35,11 +34,23 @@ describe('addon-docs enhanceSource', () => {
       });
     });
     it('transformSource', () => {
+      const transformSource = (src?: string) => (src ? `formatted: ${src}` : 'no src');
       const parameters = { ...baseContext.parameters, docs: { transformSource } };
       expect(enhanceSource({ ...baseContext, parameters }).docs.source).toEqual({
         code: 'formatted: storySource.source',
       });
     });
+    it('receives StoryContext as second argument', () => {
+      const transformSource = jest.fn();
+      const parameters = { ...baseContext.parameters, docs: { transformSource } };
+      const context = { ...baseContext, parameters };
+      enhanceSource(context);
+      expect(transformSource).toHaveBeenCalledWith(
+        baseContext.parameters.storySource.source,
+        baseContext.id,
+        context
+      );
+    });
   });
   describe('storysource source loaded w/ locationsMap', () => {
     const baseContext = {
@@ -57,11 +68,19 @@ describe('addon-docs enhanceSource', () => {
       expect(enhanceSource(baseContext)).toEqual({ docs: { source: { code: 'Source' } } });
     });
     it('transformSource', () => {
+      const transformSource = (src?: string) => (src ? `formatted: ${src}` : 'no src');
       const parameters = { ...baseContext.parameters, docs: { transformSource } };
       expect(enhanceSource({ ...baseContext, parameters }).docs.source).toEqual({
         code: 'formatted: Source',
       });
     });
+    it('receives StoryContext as second argument', () => {
+      const transformSource = jest.fn();
+      const parameters = { ...baseContext.parameters, docs: { transformSource } };
+      const context = { ...baseContext, parameters };
+      enhanceSource(context);
+      expect(transformSource).toHaveBeenCalledWith('Source', context.id, context);
+    });
   });
   describe('custom docs.source provided', () => {
     const baseContext = {
@@ -75,9 +94,21 @@ describe('addon-docs enhanceSource', () => {
       expect(enhanceSource(baseContext)).toBeNull();
     });
     it('transformSource', () => {
+      const transformSource = (src?: string) => (src ? `formatted: ${src}` : 'no src');
       const { source } = baseContext.parameters.docs;
       const parameters = { ...baseContext.parameters, docs: { source, transformSource } };
       expect(enhanceSource({ ...baseContext, parameters })).toBeNull();
     });
+    it('receives StoryContext as second argument', () => {
+      const transformSource = jest.fn();
+      const parameters = { ...baseContext.parameters, docs: { transformSource } };
+      const context = { ...baseContext, parameters };
+      enhanceSource(context);
+      expect(transformSource).toHaveBeenCalledWith(
+        context.parameters.storySource.source,
+        context.id,
+        context
+      );
+    });
   });
 });
diff --git a/addons/docs/src/blocks/enhanceSource.ts b/addons/docs/src/blocks/enhanceSource.ts
index c6e10a716e82..23e04ff83a11 100644
--- a/addons/docs/src/blocks/enhanceSource.ts
+++ b/addons/docs/src/blocks/enhanceSource.ts
@@ -54,7 +54,7 @@ export const enhanceSource = (context: StoryContext): Parameters => {
   }
 
   const input = extract(id, storySource);
-  const code = transformSource ? transformSource(input, id) : input;
+  const code = transformSource ? transformSource(input, id, context) : input;
 
   return { docs: combineParameters(docs, { source: { code } }) };
 };
diff --git a/examples/html-kitchen-sink/stories/__snapshots__/button.stories.storyshot b/examples/html-kitchen-sink/stories/__snapshots__/button.stories.storyshot
index 629692365089..1ee068c6b695 100644
--- a/examples/html-kitchen-sink/stories/__snapshots__/button.stories.storyshot
+++ b/examples/html-kitchen-sink/stories/__snapshots__/button.stories.storyshot
@@ -6,6 +6,12 @@ exports[`Storyshots Demo Button 1`] = `
 </button>
 `;
 
+exports[`Storyshots Demo Button With Custom Source Code 1`] = `
+<button>
+  Click Me
+</button>
+`;
+
 exports[`Storyshots Demo Effect 1`] = `
 <button
   id="button"
diff --git a/examples/html-kitchen-sink/stories/button.stories.js b/examples/html-kitchen-sink/stories/button.stories.js
index bf4c9ea5f998..de6d077d9e7a 100644
--- a/examples/html-kitchen-sink/stories/button.stories.js
+++ b/examples/html-kitchen-sink/stories/button.stories.js
@@ -4,19 +4,42 @@ import { useEffect } from '@storybook/client-api';
 
 export default {
   title: 'Demo',
+  argTypes: {
+    buttonText: { control: 'text', defaultValue: 'Click Me' },
+  },
 };
 
-export const Heading = () => '<h1>Hello World</h1>';
-export const Headings = () =>
-  '<h1>Hello World</h1><h2>Hello World</h2><h3>Hello World</h3><h4>Hello World</h4>';
-
-export const Button = () => {
+const Template = ({ buttonText }) => {
   const btn = document.createElement('button');
-  btn.innerHTML = 'Hello Button';
+  btn.innerHTML = buttonText;
   btn.addEventListener('click', action('Click'));
   return btn;
 };
 
+export const ButtonWithCustomSourceCode = Template.bind({});
+
+ButtonWithCustomSourceCode.args = {
+  buttonText: 'Click Me',
+};
+
+ButtonWithCustomSourceCode.parameters = {
+  docs: {
+    transformSource: (src, id, context) => {
+      return Template(context.args).outerHTML;
+    },
+  },
+};
+
+export const Button = Template.bind({});
+
+Button.args = {
+  buttonText: 'Click Me',
+};
+
+export const Heading = () => '<h1>Hello World</h1>';
+export const Headings = () =>
+  '<h1>Hello World</h1><h2>Hello World</h2><h3>Hello World</h3><h4>Hello World</h4>';
+
 export const Effect = () => {
   useEffect(() => {
     document.getElementById('button').style.backgroundColor = 'yellow';
diff --git a/lib/api/src/version.ts b/lib/api/src/version.ts
index 5e9575a55427..550fc7ad3606 100644
--- a/lib/api/src/version.ts
+++ b/lib/api/src/version.ts
@@ -1 +1 @@
-export const version = '6.1.0-alpha.1';
\ No newline at end of file
+export const version = '6.1.0-alpha.1';
diff --git a/lib/cli/versions.json b/lib/cli/versions.json
index b7543a709ee2..c56efdcad6db 100644
--- a/lib/cli/versions.json
+++ b/lib/cli/versions.json
@@ -52,4 +52,4 @@
   "@storybook/ui": "6.1.0-alpha.1",
   "@storybook/vue": "6.1.0-alpha.1",
   "@storybook/web-components": "6.1.0-alpha.1"
-}
\ No newline at end of file
+}

From 1457d269443098d512de042aee1cb9c08dcd64e0 Mon Sep 17 00:00:00 2001
From: Nikki Reece <nicole.reece@monsanto.com>
Date: Mon, 31 Aug 2020 11:03:05 -0500
Subject: [PATCH 2/3] update snapshots

---
 .../stories/__snapshots__/button.stories.storyshot              | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/html-kitchen-sink/stories/__snapshots__/button.stories.storyshot b/examples/html-kitchen-sink/stories/__snapshots__/button.stories.storyshot
index 1ee068c6b695..0be594ce8b22 100644
--- a/examples/html-kitchen-sink/stories/__snapshots__/button.stories.storyshot
+++ b/examples/html-kitchen-sink/stories/__snapshots__/button.stories.storyshot
@@ -2,7 +2,7 @@
 
 exports[`Storyshots Demo Button 1`] = `
 <button>
-  Hello Button
+  Click Me
 </button>
 `;
 

From 05735b6e87b8fdd257ace0815a4d1cd9f999c532 Mon Sep 17 00:00:00 2001
From: Nikki Reece <nicole.reece@monsanto.com>
Date: Tue, 1 Sep 2020 08:17:38 -0500
Subject: [PATCH 3/3] test with one test instead of three

---
 addons/docs/src/blocks/enhanceSource.test.ts | 18 ------------------
 1 file changed, 18 deletions(-)

diff --git a/addons/docs/src/blocks/enhanceSource.test.ts b/addons/docs/src/blocks/enhanceSource.test.ts
index 695e66abad1d..7894c5cd6bb2 100644
--- a/addons/docs/src/blocks/enhanceSource.test.ts
+++ b/addons/docs/src/blocks/enhanceSource.test.ts
@@ -74,13 +74,6 @@ describe('addon-docs enhanceSource', () => {
         code: 'formatted: Source',
       });
     });
-    it('receives StoryContext as second argument', () => {
-      const transformSource = jest.fn();
-      const parameters = { ...baseContext.parameters, docs: { transformSource } };
-      const context = { ...baseContext, parameters };
-      enhanceSource(context);
-      expect(transformSource).toHaveBeenCalledWith('Source', context.id, context);
-    });
   });
   describe('custom docs.source provided', () => {
     const baseContext = {
@@ -99,16 +92,5 @@ describe('addon-docs enhanceSource', () => {
       const parameters = { ...baseContext.parameters, docs: { source, transformSource } };
       expect(enhanceSource({ ...baseContext, parameters })).toBeNull();
     });
-    it('receives StoryContext as second argument', () => {
-      const transformSource = jest.fn();
-      const parameters = { ...baseContext.parameters, docs: { transformSource } };
-      const context = { ...baseContext, parameters };
-      enhanceSource(context);
-      expect(transformSource).toHaveBeenCalledWith(
-        context.parameters.storySource.source,
-        context.id,
-        context
-      );
-    });
   });
 });