Docs: Fix attachment logic

code/lib/preview-api/src/modules/preview-web/PreviewWeb.test.ts

@@ -642,7 +642,7 @@ describe('PreviewWeb', () => {
       });
     });
 
-    describe('template docs entries', () => {
+    describe('CSF docs entries', () => {
       it('always renders in docs viewMode', async () => {
         document.location.search = '?id=component-one--docs';
         await createAndRenderPreview();
@@ -679,7 +679,7 @@ describe('PreviewWeb', () => {
         expect(importFn).toHaveBeenCalledWith('./src/ExtraComponentOne.stories.js');
       });
 
-      it('renders with componentStories loaded from both story files', async () => {
+      it('renders with componentStories loaded from the attached CSF file', async () => {
         document.location.search = '?id=component-one--docs&viewMode=docs';
         await createAndRenderPreview();
 

code/lib/preview-api/src/modules/preview-web/docs-context/DocsContext.test.ts

@@ -1,31 +1,15 @@
 import { Channel } from '@storybook/channels';
-import type { Renderer, CSFFile, PreparedStory } from '@storybook/types';
+import type { Renderer } from '@storybook/types';
 import type { StoryStore } from '../../store';
 
 import { DocsContext } from './DocsContext';
+import { csfFileParts } from './test-utils';
 
 const channel = new Channel();
 const renderStoryToElement = jest.fn();
 
 describe('resolveModuleExport', () => {
-  // These compose the raw exports of the CSF file
-  const component = {};
-  const metaExport = { component };
-  const storyExport = {};
-  const moduleExports = { default: metaExport, story: storyExport };
-
-  // This is the prepared story + CSF file after SB has processed them
-  const storyAnnotations = {
-    id: 'meta--story',
-    moduleExport: storyExport,
-  } as CSFFile['stories'][string];
-  const story = { id: 'meta--story', moduleExport: storyExport } as PreparedStory;
-  const meta = { id: 'meta', title: 'Meta', component, moduleExports } as CSFFile['meta'];
-  const csfFile = {
-    stories: { story: storyAnnotations },
-    meta,
-    moduleExports,
-  } as CSFFile;
+  const { story, csfFile, storyExport, metaExport, moduleExports, component } = csfFileParts();
 
   const store = {
     componentStoriesFromCSFFile: () => [story],

code/lib/preview-api/src/modules/preview-web/docs-context/DocsContext.ts

@@ -31,8 +31,7 @@ export class DocsContext<TRenderer extends Renderer> implements DocsContextProps
     protected store: StoryStore<TRenderer>,
     public renderStoryToElement: DocsContextProps['renderStoryToElement'],
     /** The CSF files known (via the index) to be refererenced by this docs file */
-    csfFiles: CSFFile<TRenderer>[],
-    componentStoriesFromAllCsfFiles = true
+    csfFiles: CSFFile<TRenderer>[]
   ) {
     this.storyIdToCSFFile = new Map();
     this.exportToStory = new Map();
@@ -41,14 +40,14 @@ export class DocsContext<TRenderer extends Renderer> implements DocsContextProps
     this.componentStoriesValue = [];
 
     csfFiles.forEach((csfFile, index) => {
-      this.referenceCSFFile(csfFile, componentStoriesFromAllCsfFiles || index === 0);
+      this.referenceCSFFile(csfFile);
     });
   }
 
   // This docs entry references this CSF file and can syncronously load the stories, as well
   // as reference them by module export. If the CSF is part of the "component" stories, they
   // can also be referenced by name and are in the componentStories list.
-  referenceCSFFile(csfFile: CSFFile<TRenderer>, addToComponentStories: boolean) {
+  referenceCSFFile(csfFile: CSFFile<TRenderer>) {
     this.exportsToCSFFile.set(csfFile.moduleExports, csfFile);
     // Also set the default export as the component's exports,
     // to allow `import ButtonStories from './Button.stories'`
@@ -63,18 +62,26 @@ export class DocsContext<TRenderer extends Renderer> implements DocsContextProps
         throw new Error(`Unexpected missing story ${annotation.id} from referenced CSF file.`);
       this.exportToStory.set(annotation.moduleExport, story);
     });
+  }
 
-    if (addToComponentStories) {
-      stories.forEach((story) => {
-        this.nameToStoryId.set(story.name, story.id);
-        this.componentStoriesValue.push(story);
-        if (!this.primaryStory) this.primaryStory = story;
-      });
+  attachCSFFile(csfFile: CSFFile<TRenderer>) {
+    if (!this.exportsToCSFFile.has(csfFile.moduleExports)) {
+      throw new Error('Cannot attach a CSF file that has not been referenced');
     }
+
+    const stories = this.store.componentStoriesFromCSFFile({ csfFile });
+    stories.forEach((story) => {
+      this.nameToStoryId.set(story.name, story.id);
+      this.componentStoriesValue.push(story);
+      if (!this.primaryStory) this.primaryStory = story;
+    });
   }
 
   setMeta(metaExports: ModuleExports) {
-    // Do nothing (this is really only used by external docs)
+    const resolved = this.resolveModuleExport(metaExports);
+    if (resolved.type !== 'meta')
+      throw new Error('Cannot reference a non-meta or module export in <Meta of={} />');
+    this.attachCSFFile(resolved.csfFile);
   }
 
   resolveModuleExport(moduleExport: ModuleExport, metaExports?: ModuleExports) {

code/lib/preview-api/src/modules/preview-web/docs-context/test-utils.ts

@@ -0,0 +1,33 @@
+import type { CSFFile, PreparedStory } from '@storybook/types';
+
+export function csfFileParts() {
+  // These compose the raw exports of the CSF file
+  const component = {};
+  const metaExport = { component };
+  const storyExport = {};
+  const moduleExports = { default: metaExport, story: storyExport };
+
+  // This is the prepared story + CSF file after SB has processed them
+  const storyAnnotations = {
+    id: 'meta--story',
+    moduleExport: storyExport,
+  } as CSFFile['stories'][string];
+  const story = { id: 'meta--story', moduleExport: storyExport } as PreparedStory;
+  const meta = { id: 'meta', title: 'Meta', component, moduleExports } as CSFFile['meta'];
+  const csfFile = {
+    stories: { story: storyAnnotations },
+    meta,
+    moduleExports,
+  } as CSFFile;
+
+  return {
+    component,
+    metaExport,
+    storyExport,
+    moduleExports,
+    storyAnnotations,
+    story,
+    meta,
+    csfFile,
+  };
+}

code/lib/preview-api/src/modules/preview-web/render/CsfDocsRender.test.ts

@@ -1,9 +1,16 @@
 import { Channel } from '@storybook/channels';
-import type { Renderer, DocsIndexEntry } from '@storybook/types';
+import type {
+  Renderer,
+  DocsIndexEntry,
+  PreparedStory,
+  ModuleExports,
+  CSFFile,
+} from '@storybook/types';
 import type { StoryStore } from '../../store';
 import { PREPARE_ABORTED } from './Render';
 
 import { CsfDocsRender } from './CsfDocsRender';
+import { csfFileParts } from '../docs-context/test-utils';
 
 const entry = {
   type: 'docs',
@@ -23,28 +30,47 @@ const createGate = (): [Promise<any | undefined>, (_?: any) => void] => {
   return [gate, openGate];
 };
 
-describe('CsfDocsRender', () => {
-  it('throws PREPARE_ABORTED if torndown during prepare', async () => {
-    const [importGate, openImportGate] = createGate();
-    const mockStore = {
-      loadEntry: jest.fn(async () => {
-        await importGate;
-        return {};
-      }),
-    };
+it('throws PREPARE_ABORTED if torndown during prepare', async () => {
+  const [importGate, openImportGate] = createGate();
+  const mockStore = {
+    loadEntry: jest.fn(async () => {
+      await importGate;
+      return {};
+    }),
+  };
 
-    const render = new CsfDocsRender(
-      new Channel(),
-      mockStore as unknown as StoryStore<Renderer>,
-      entry
-    );
+  const render = new CsfDocsRender(
+    new Channel(),
+    mockStore as unknown as StoryStore<Renderer>,
+    entry
+  );
 
-    const preparePromise = render.prepare();
+  const preparePromise = render.prepare();
 
-    render.teardown();
+  render.teardown();
 
-    openImportGate();
+  openImportGate();
 
-    await expect(preparePromise).rejects.toThrowError(PREPARE_ABORTED);
-  });
+  await expect(preparePromise).rejects.toThrowError(PREPARE_ABORTED);
+});
+
+it('attached immediately', async () => {
+  const { story, csfFile, moduleExports } = csfFileParts();
+
+  const store = {
+    loadEntry: () => ({
+      entryExports: moduleExports,
+      csfFiles: [],
+    }),
+    processCSFFileWithCache: () => csfFile,
+    componentStoriesFromCSFFile: () => [story],
+    storyFromCSFFile: () => story,
+  } as unknown as StoryStore<Renderer>;
+
+  const render = new CsfDocsRender(new Channel(), store, entry);
+  await render.prepare();
+
+  const context = render.docsContext(jest.fn());
+
+  expect(context.storyById()).toEqual(story);
 });

code/lib/preview-api/src/modules/preview-web/render/CsfDocsRender.ts

@@ -85,19 +85,28 @@ export class CsfDocsRender<TRenderer extends Renderer> implements Render<TRender
     );
   }
 
+  docsContext(renderStoryToElement: DocsContextProps['renderStoryToElement']) {
+    if (!this.csfFiles) throw new Error('Cannot render docs before preparing');
+    const docsContext = new DocsContext<TRenderer>(
+      this.channel,
+      this.store,
+      renderStoryToElement,
+      this.csfFiles
+    );
+    // All referenced CSF files should be attached for CSF docs
+    //  - When you create two CSF files that both reference the same title, they are combined into
+    //    a single CSF docs entry with a `storiesImport` defined.
+    this.csfFiles.forEach((csfFile) => docsContext.attachCSFFile(csfFile));
+    return docsContext;
+  }
+
   async renderToElement(
     canvasElement: TRenderer['canvasElement'],
     renderStoryToElement: DocsContextProps['renderStoryToElement']
   ) {
     if (!this.story || !this.csfFiles) throw new Error('Cannot render docs before preparing');
 
-    const docsContext = new DocsContext<TRenderer>(
-      this.channel,
-      this.store,
-      renderStoryToElement,
-      this.csfFiles,
-      true
-    );
+    const docsContext = this.docsContext(renderStoryToElement);
 
     const { docs: docsParameter } = this.story.parameters || {};
 

code/lib/preview-api/src/modules/preview-web/render/MdxDocsRender.test.ts

@@ -4,6 +4,7 @@ import type { StoryStore } from '../../store';
 import { PREPARE_ABORTED } from './Render';
 
 import { MdxDocsRender } from './MdxDocsRender';
+import { csfFileParts } from '../docs-context/test-utils';
 
 const entry = {
   type: 'docs',
@@ -22,28 +23,58 @@ const createGate = (): [Promise<any | undefined>, (_?: any) => void] => {
   return [gate, openGate];
 };
 
-describe('MdxDocsRender', () => {
-  it('throws PREPARE_ABORTED if torndown during prepare', async () => {
-    const [importGate, openImportGate] = createGate();
-    const mockStore = {
-      loadEntry: jest.fn(async () => {
-        await importGate;
-        return {};
-      }),
-    };
+it('throws PREPARE_ABORTED if torndown during prepare', async () => {
+  const [importGate, openImportGate] = createGate();
+  const mockStore = {
+    loadEntry: jest.fn(async () => {
+      await importGate;
+      return {};
+    }),
+  };
 
-    const render = new MdxDocsRender(
-      new Channel(),
-      mockStore as unknown as StoryStore<Renderer>,
-      entry
-    );
+  const render = new MdxDocsRender(
+    new Channel(),
+    mockStore as unknown as StoryStore<Renderer>,
+    entry
+  );
 
-    const preparePromise = render.prepare();
+  const preparePromise = render.prepare();
 
-    render.teardown();
+  render.teardown();
 
-    openImportGate();
+  openImportGate();
 
-    await expect(preparePromise).rejects.toThrowError(PREPARE_ABORTED);
+  await expect(preparePromise).rejects.toThrowError(PREPARE_ABORTED);
+});
+
+describe('attaching', () => {
+  const { story, csfFile, moduleExports } = csfFileParts();
+  const store = {
+    loadEntry: () => ({
+      entryExports: moduleExports,
+      csfFiles: [csfFile],
+    }),
+    processCSFFileWithCache: () => csfFile,
+    componentStoriesFromCSFFile: () => [story],
+    storyFromCSFFile: () => story,
+  } as unknown as StoryStore<Renderer>;
+
+  it('is not attached if you do not call setMeta', async () => {
+    const render = new MdxDocsRender(new Channel(), store, entry);
+    await render.prepare();
+
+    const context = render.docsContext(jest.fn());
+
+    expect(context.storyById).toThrow('No primary story defined');
+  });
+
+  it('is attached if you call setMeta', async () => {
+    const render = new MdxDocsRender(new Channel(), store, entry);
+    await render.prepare();
+
+    const context = render.docsContext(jest.fn());
+    context.setMeta(moduleExports);
+
+    expect(context.storyById()).toEqual(story);
   });
 });

code/lib/preview-api/src/modules/preview-web/render/MdxDocsRender.ts

@@ -69,23 +69,29 @@ export class MdxDocsRender<TRenderer extends Renderer> implements Render<TRender
     );
   }
 
+  docsContext(renderStoryToElement: DocsContextProps['renderStoryToElement']) {
+    if (!this.csfFiles) throw new Error('Cannot render docs before preparing');
+
+    // NOTE we do *not* attach any CSF file yet. We wait for `setMeta`
+    // ie the CSF file is attached via `<Meta of={} />`
+    return new DocsContext<TRenderer>(
+      this.channel,
+      this.store,
+      renderStoryToElement,
+      this.csfFiles
+    );
+  }
+
   async renderToElement(
     canvasElement: TRenderer['canvasElement'],
     renderStoryToElement: DocsContextProps['renderStoryToElement']
   ) {
     if (!this.exports || !this.csfFiles || !this.store.projectAnnotations)
       throw new Error('Cannot render docs before preparing');
 
-    const docsContext = new DocsContext<TRenderer>(
-      this.channel,
-      this.store,
-      renderStoryToElement,
-      this.csfFiles,
-      false
-    );
+    const docsContext = this.docsContext(renderStoryToElement);
 
     const { docs } = this.store.projectAnnotations.parameters || {};
-
     if (!docs)
       throw new Error(
         `Cannot render a story in viewMode=docs if \`@storybook/addon-docs\` is not installed`