Addon-docs: Change 2nd argument of transformSource to the storyContext

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;
     },

addons/docs/src/blocks/enhanceSource.test.ts

@@ -11,35 +11,43 @@ const emptyContext: StoryContext = {
   parameters: {},
 };
 
-const transformSource = (src?: string) => (src ? `formatted: ${src}` : 'no src');
-
 describe('addon-docs enhanceSource', () => {
   describe('no source loaded', () => {
     const baseContext = emptyContext;
     it('no transformSource', () => {
       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({
         docs: { source: { code: 'storySource.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: 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();

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 } }) };
 };