From 4664dcc2eb6993fca4fd467b26a639f4b52faa9c Mon Sep 17 00:00:00 2001
From: Blaine Bublitz <blaine.bublitz@gmail.com>
Date: Wed, 26 Aug 2020 11:41:25 -0700
Subject: [PATCH] Addon-docs: Change 2nd argument of transformSource to the
 storyContext

---
 addons/docs/docs/recipes.md                  |  2 +-
 addons/docs/src/blocks/enhanceSource.test.ts | 23 +++++++++++++++++---
 addons/docs/src/blocks/enhanceSource.ts      |  2 +-
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/addons/docs/docs/recipes.md b/addons/docs/docs/recipes.md
index e08fe60c0f0e..7fd560e95114 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, 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..006463374fd0 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,14 +18,16 @@ 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();
     });
   });
   describe('custom/mdx source loaded', () => {
+    const source = 'storySource.source';
     const baseContext = {
       ...emptyContext,
-      parameters: { storySource: { source: 'storySource.source' } },
+      parameters: { storySource: { source } },
     };
     it('no transformSource', () => {
       expect(enhanceSource(baseContext)).toEqual({
@@ -35,11 +35,19 @@ 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 context as 2nd argument', () => {
+      const transformSource = jest.fn();
+      const parameters = { ...baseContext.parameters, docs: { transformSource } };
+      const context = { ...baseContext, parameters };
+      enhanceSource(context);
+      expect(transformSource).toHaveBeenCalledWith(source, context);
+    });
   });
   describe('storysource source loaded w/ locationsMap', () => {
     const baseContext = {
@@ -57,11 +65,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 context as 2nd argument', () => {
+      const transformSource = jest.fn();
+      const parameters = { ...baseContext.parameters, docs: { transformSource } };
+      const context = { ...baseContext, parameters };
+      enhanceSource(context);
+      expect(transformSource).toHaveBeenCalledWith('Source', context);
+    });
   });
   describe('custom docs.source provided', () => {
     const baseContext = {
@@ -75,6 +91,7 @@ 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();
diff --git a/addons/docs/src/blocks/enhanceSource.ts b/addons/docs/src/blocks/enhanceSource.ts
index c6e10a716e82..252155205e34 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, context) : input;
 
   return { docs: combineParameters(docs, { source: { code } }) };
 };