Refactor PreviewWeb rendering into Story/DocsRender

addons/docs/src/blocks/Story.tsx

@@ -130,25 +130,8 @@ const Story: FunctionComponent<StoryProps> = (props) => {
   useEffect(() => {
     let cleanup: () => void;
     if (story && storyRef.current) {
-      const { componentId, id, title, name } = story;
-      const renderContext = {
-        componentId,
-        title,
-        kind: title,
-        id,
-        name,
-        story: name,
-        // TODO what to do when these fail?
-        showMain: () => {},
-        showError: () => {},
-        showException: () => {},
-      };
-      cleanup = context.renderStoryToElement({
-        story,
-        renderContext,
-        element: storyRef.current as HTMLElement,
-        viewMode: 'docs',
-      });
+      const element = storyRef.current as HTMLElement;
+      cleanup = context.renderStoryToElement(story, element);
       setShowLoader(false);
     }
     return () => cleanup && cleanup();

lib/preview-web/src/DocsRender.ts

@@ -0,0 +1,90 @@
+import global from 'global';
+import {
+  AnyFramework,
+  StoryId,
+  ViewMode,
+  StoryContextForLoaders,
+  StoryContext,
+} from '@storybook/csf';
+import { Story, StoryStore, CSFFile } from '@storybook/store';
+import { Channel } from '@storybook/addons';
+import { DOCS_RENDERED } from '@storybook/core-events';
+
+import { DocsContextProps } from './types';
+
+export class DocsRender<CanvasElement extends HTMLElement | void, TFramework extends AnyFramework> {
+  public story?: Story<TFramework>;
+
+  private canvasElement?: CanvasElement;
+
+  private context?: DocsContextProps;
+
+  public disableKeyListeners = false;
+
+  constructor(
+    private channel: Channel,
+    private store: StoryStore<TFramework>,
+    public id: StoryId,
+    story?: Story<TFramework>
+  ) {
+    if (story) this.story = story;
+  }
+
+  async prepare() {
+    this.story = await this.store.loadStory({ storyId: this.id });
+  }
+
+  async renderToElement(
+    canvasElement: CanvasElement,
+    renderStoryToElement: DocsContextProps['renderStoryToElement']
+  ) {
+    this.canvasElement = canvasElement;
+
+    const { id, title, name } = this.story;
+    const csfFile: CSFFile<TFramework> = await this.store.loadCSFFileByStoryId(this.id);
+
+    this.context = {
+      id,
+      title,
+      name,
+      // NOTE: these two functions are *sync* so cannot access stories from other CSF files
+      storyById: (storyId: StoryId) => this.store.storyFromCSFFile({ storyId, csfFile }),
+      componentStories: () => this.store.componentStoriesFromCSFFile({ csfFile }),
+      loadStory: (storyId: StoryId) => this.store.loadStory({ storyId }),
+      renderStoryToElement: renderStoryToElement.bind(this),
+      getStoryContext: (renderedStory: Story<TFramework>) =>
+        ({
+          ...this.store.getStoryContext(renderedStory),
+          viewMode: 'docs' as ViewMode,
+        } as StoryContextForLoaders<TFramework>),
+      // Put all the storyContext fields onto the docs context for back-compat
+      ...(!global.FEATURES?.breakingChangesV7 && this.store.getStoryContext(this.story)),
+    };
+
+    return this.render();
+  }
+
+  async render() {
+    if (!this.story || !this.context || !this.canvasElement)
+      throw new Error('DocsRender not ready to render');
+
+    const renderer = await import('./renderDocs');
+    renderer.renderDocs(this.story, this.context, this.canvasElement, () =>
+      this.channel.emit(DOCS_RENDERED, this.id)
+    );
+  }
+
+  async rerender() {
+    // NOTE: in modern inline render mode, each story is rendered via
+    // `preview.renderStoryToElement` which means the story will track
+    // its own re-renders. Thus there will be no need to re-render the whole
+    // docs page when a single story changes.
+    if (!global.FEATURES?.modernInlineRender) await this.render();
+  }
+
+  async teardown({ viewModeChanged }: { viewModeChanged?: boolean } = {}) {
+    if (!viewModeChanged || !this.canvasElement) return;
+    const renderer = await import('./renderDocs');
+    renderer.unmountDocs(this.canvasElement);
+  }
+}

lib/preview-web/src/PreviewWeb.test.ts

@@ -79,6 +79,9 @@ async function createAndRenderPreview({
   getProjectAnnotations?: () => WebProjectAnnotations<AnyFramework>;
 } = {}) {
   const preview = new PreviewWeb();
+  (
+    preview.view.prepareForDocs as jest.MockedFunction<typeof preview.view.prepareForDocs>
+  ).mockReturnValue('docs-element' as any);
   await preview.initialize({
     importFn: inputImportFn,
     getProjectAnnotations: inputGetProjectAnnotations,
@@ -595,7 +598,7 @@ describe('PreviewWeb', () => {
               }),
             }),
           }),
-          undefined,
+          'docs-element',
           expect.any(Function)
         );
       });
@@ -617,6 +620,7 @@ describe('PreviewWeb', () => {
 
       emitter.emit(Events.UPDATE_GLOBALS, { globals: { foo: 'bar' } });
 
+      await waitForEvents([Events.GLOBALS_UPDATED]);
       expect(mockChannel.emit).toHaveBeenCalledWith(Events.GLOBALS_UPDATED, {
         globals: { a: 'b', foo: 'bar' },
         initialGlobals: { a: 'b' },
@@ -688,6 +692,7 @@ describe('PreviewWeb', () => {
         updatedArgs: { new: 'arg' },
       });
 
+      await waitForEvents([Events.STORY_ARGS_UPDATED]);
       expect(mockChannel.emit).toHaveBeenCalledWith(Events.STORY_ARGS_UPDATED, {
         storyId: 'component-one--a',
         args: { foo: 'a', new: 'arg' },
@@ -935,20 +940,85 @@ describe('PreviewWeb', () => {
       });
     });
 
-    describe('in docs mode', () => {
+    describe('in docs mode, old inline render', () => {
       it('re-renders the docs container', async () => {
         document.location.search = '?id=component-one--a&viewMode=docs';
 
         await createAndRenderPreview();
 
+        (ReactDOM.render as jest.MockedFunction<typeof ReactDOM.render>).mockClear();
         mockChannel.emit.mockClear();
         emitter.emit(Events.UPDATE_STORY_ARGS, {
           storyId: 'component-one--a',
           updatedArgs: { new: 'arg' },
         });
         await waitForRender();
 
-        expect(ReactDOM.render).toHaveBeenCalledTimes(2);
+        expect(ReactDOM.render).toHaveBeenCalledTimes(1);
+      });
+    });
+
+    describe('in docs mode, modern inline render', () => {
+      beforeEach(() => {
+        global.FEATURES.modernInlineRender = true;
+      });
+      afterEach(() => {
+        global.FEATURES.modernInlineRender = true;
+      });
+      it('does not re-render the docs container', async () => {
+        document.location.search = '?id=component-one--a&viewMode=docs';
+
+        await createAndRenderPreview();
+
+        (ReactDOM.render as jest.MockedFunction<typeof ReactDOM.render>).mockClear();
+        mockChannel.emit.mockClear();
+        emitter.emit(Events.UPDATE_STORY_ARGS, {
+          storyId: 'component-one--a',
+          updatedArgs: { new: 'arg' },
+        });
+        await waitForEvents([Events.STORY_ARGS_UPDATED]);
+
+        expect(ReactDOM.render).not.toHaveBeenCalled();
+      });
+
+      describe('when renderStoryToElement was called', () => {
+        it('re-renders the story', async () => {
+          document.location.search = '?id=component-one--a&viewMode=docs';
+
+          const preview = await createAndRenderPreview();
+          await waitForRender();
+
+          mockChannel.emit.mockClear();
+          const story = await preview.storyStore.loadStory({ storyId: 'component-one--a' });
+          preview.renderStoryToElement(story, 'story-element' as any);
+          await waitForRender();
+
+          expect(projectAnnotations.renderToDOM).toHaveBeenCalledWith(
+            expect.objectContaining({
+              storyContext: expect.objectContaining({
+                args: { foo: 'a' },
+              }),
+            }),
+            'story-element'
+          );
+
+          (ReactDOM.render as jest.MockedFunction<typeof ReactDOM.render>).mockClear();
+          mockChannel.emit.mockClear();
+          emitter.emit(Events.UPDATE_STORY_ARGS, {
+            storyId: 'component-one--a',
+            updatedArgs: { new: 'arg' },
+          });
+          await waitForRender();
+
+          expect(projectAnnotations.renderToDOM).toHaveBeenCalledWith(
+            expect.objectContaining({
+              storyContext: expect.objectContaining({
+                args: { foo: 'a', new: 'arg' },
+              }),
+            }),
+            'story-element'
+          );
+        });
       });
     });
   });
@@ -964,6 +1034,7 @@ describe('PreviewWeb', () => {
         updatedArgs: { foo: 'new' },
       });
 
+      await waitForEvents([Events.STORY_ARGS_UPDATED]);
       expect(mockChannel.emit).toHaveBeenCalledWith(Events.STORY_ARGS_UPDATED, {
         storyId: 'component-one--a',
         args: { foo: 'new' },
@@ -992,6 +1063,7 @@ describe('PreviewWeb', () => {
         storyId: 'component-one--a',
         updatedArgs: { foo: 'new', new: 'value' },
       });
+      await waitForEvents([Events.STORY_ARGS_UPDATED]);
 
       mockChannel.emit.mockClear();
       emitter.emit(Events.RESET_STORY_ARGS, {
@@ -1012,6 +1084,7 @@ describe('PreviewWeb', () => {
         undefined // this is coming from view.prepareForStory, not super important
       );
 
+      await waitForEvents([Events.STORY_ARGS_UPDATED]);
       expect(mockChannel.emit).toHaveBeenCalledWith(Events.STORY_ARGS_UPDATED, {
         storyId: 'component-one--a',
         args: { foo: 'a', new: 'value' },
@@ -1026,6 +1099,7 @@ describe('PreviewWeb', () => {
         storyId: 'component-one--a',
         updatedArgs: { foo: 'new', new: 'value' },
       });
+      await waitForEvents([Events.STORY_ARGS_UPDATED]);
 
       mockChannel.emit.mockClear();
       emitter.emit(Events.RESET_STORY_ARGS, {
@@ -1044,6 +1118,8 @@ describe('PreviewWeb', () => {
         }),
         undefined // this is coming from view.prepareForStory, not super important
       );
+
+      await waitForEvents([Events.STORY_ARGS_UPDATED]);
       expect(mockChannel.emit).toHaveBeenCalledWith(Events.STORY_ARGS_UPDATED, {
         storyId: 'component-one--a',
         args: { foo: 'a' },
@@ -1766,7 +1842,7 @@ describe('PreviewWeb', () => {
               }),
             }),
           }),
-          undefined,
+          'docs-element',
           expect.any(Function)
         );
       });