Merge pull request #18457 from storybookjs/tom/sb-415-return-cleanup-function-from-rendertodom

Core: Allow a teardown function to be returned from `renderToDOM`

app/react/src/client/preview/render.tsx

@@ -147,4 +147,6 @@ export async function renderToDOM(
   }
 
   await renderElement(element, domElement);
+
+  return () => unmountElement(domElement);
 }

examples/react-ts/.storybook/main.ts

@@ -37,7 +37,7 @@ const config: StorybookConfig = {
   },
   features: {
     postcss: false,
-    // modernInlineRender: true,
+    modernInlineRender: true,
     storyStoreV7: !global.navigator?.userAgent?.match?.('jsdom'),
     buildStoriesJson: true,
     babelModeV7: true,

examples/react-ts/src/__snapshots__/storyshots.test.ts.snap

@@ -4041,6 +4041,51 @@ exports[`Storyshots Demo/Examples / Emoji Button With Args 1`] = `
 </button>
 `;
 
+exports[`Storyshots Demo/button2 One 1`] = `
+<button
+  type="button"
+>
+   
+  one
+</button>
+`;
+
+exports[`Storyshots Demo/button2 Three 1`] = `
+<button
+  type="button"
+>
+   
+  three
+</button>
+`;
+
+exports[`Storyshots Demo/button2 Two 1`] = `
+<button
+  type="button"
+>
+   
+  two
+</button>
+`;
+
+exports[`Storyshots Demo/button3 Five 1`] = `
+<button
+  type="button"
+>
+   
+  five
+</button>
+`;
+
+exports[`Storyshots Demo/button3 Four 1`] = `
+<button
+  type="button"
+>
+   
+  four
+</button>
+`;
+
 exports[`Storyshots Docs/ButtonMdx Basic 1`] = `
 <button
   type="button"

examples/react-ts/src/button.tsx

@@ -1,4 +1,4 @@
-import React, { ComponentType, ButtonHTMLAttributes } from 'react';
+import React, { ComponentType, ButtonHTMLAttributes, useEffect } from 'react';
 
 export interface ButtonProps extends ButtonHTMLAttributes<HTMLButtonElement> {
   /**
@@ -12,8 +12,15 @@ export interface ButtonProps extends ButtonHTMLAttributes<HTMLButtonElement> {
   icon?: ComponentType;
 }
 
-export const Button = ({ label = 'Hello', icon: Icon, ...props }: ButtonProps) => (
-  <button type="button" {...props}>
-    {Icon ? <Icon /> : null} {label}
-  </button>
-);
+export const Button = ({ label = 'Hello', icon: Icon, ...props }: ButtonProps) => {
+  useEffect(() => {
+    const fn = () => console.log(`click ${label}`);
+    global.window.document.querySelector('body')?.addEventListener('click', fn);
+    return () => global.window.document.querySelector('body')?.removeEventListener('click', fn);
+  });
+  return (
+    <button type="button" {...props}>
+      {Icon ? <Icon /> : null} {label}
+    </button>
+  );
+};

examples/react-ts/src/button2.stories.tsx

@@ -0,0 +1,10 @@
+import { Button } from './button';
+
+export default {
+  component: Button,
+  title: 'button2',
+};
+
+export const one = { args: { label: 'one' } };
+export const two = { args: { label: 'two' } };
+export const three = { args: { label: 'three' } };

examples/react-ts/src/button3.stories.tsx

@@ -0,0 +1,9 @@
+import { Button } from './button';
+
+export default {
+  component: Button,
+  title: 'button3',
+};
+
+export const four = { args: { label: 'four' } };
+export const five = { args: { label: 'five' } };

lib/preview-web/src/Preview.tsx

@@ -23,6 +23,7 @@ import {
   StoryIndex,
   PromiseLike,
   WebProjectAnnotations,
+  RenderToDOM,
 } from '@storybook/store';
 
 import { StoryRender } from './StoryRender';
@@ -45,7 +46,7 @@ export class Preview<TFramework extends AnyFramework> {
 
   importFn?: ModuleImportFn;
 
-  renderToDOM: WebProjectAnnotations<TFramework>['renderToDOM'];
+  renderToDOM: RenderToDOM<TFramework>;
 
   storyRenders: StoryRender<TFramework>[] = [];
 

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

@@ -7,7 +7,7 @@ import {
   STORY_RENDER_PHASE_CHANGED,
   STORY_THREW_EXCEPTION,
 } from '@storybook/core-events';
-import { StoryIndex } from '@storybook/store';
+import { StoryIndex, TeardownRenderToDOM } from '@storybook/store';
 import { RenderPhase } from './PreviewWeb';
 
 export const componentOneExports = {
@@ -32,12 +32,13 @@ export const importFn = jest.fn(async (path) => {
   return path === './src/ComponentOne.stories.js' ? componentOneExports : componentTwoExports;
 });
 
+export const teardownRenderToDOM: jest.Mock<TeardownRenderToDOM> = jest.fn();
 export const projectAnnotations = {
   globals: { a: 'b' },
   globalTypes: {},
   decorators: [jest.fn((s) => s())],
   render: jest.fn(),
-  renderToDOM: jest.fn(),
+  renderToDOM: jest.fn().mockReturnValue(teardownRenderToDOM),
 };
 export const getProjectAnnotations = () => projectAnnotations;
 

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

@@ -44,6 +44,7 @@ import {
   waitForRender,
   waitForQuiescence,
   waitForRenderPhase,
+  teardownRenderToDOM,
 } from './PreviewWeb.mockdata';
 
 jest.mock('./WebView');
@@ -120,7 +121,8 @@ beforeEach(() => {
   componentOneExports.default.loaders[0].mockReset().mockImplementation(async () => ({ l: 7 }));
   componentOneExports.default.parameters.docs.container.mockClear();
   componentOneExports.a.play.mockReset();
-  projectAnnotations.renderToDOM.mockReset();
+  teardownRenderToDOM.mockReset();
+  projectAnnotations.renderToDOM.mockReset().mockReturnValue(teardownRenderToDOM);
   projectAnnotations.render.mockClear();
   projectAnnotations.decorators[0].mockClear();
   (ReactDOM.render as any as jest.Mock<typeof ReactDOM.render>)
@@ -470,7 +472,7 @@ describe('PreviewWeb', () => {
 
       it('renders exception if renderToDOM throws', async () => {
         const error = new Error('error');
-        projectAnnotations.renderToDOM.mockImplementationOnce(() => {
+        projectAnnotations.renderToDOM.mockImplementation(() => {
           throw error;
         });
 
@@ -519,9 +521,7 @@ describe('PreviewWeb', () => {
 
       it('renders error if the story calls showError', async () => {
         const error = { title: 'title', description: 'description' };
-        projectAnnotations.renderToDOM.mockImplementationOnce((context) =>
-          context.showError(error)
-        );
+        projectAnnotations.renderToDOM.mockImplementation((context) => context.showError(error));
 
         document.location.search = '?id=component-one--a';
         const preview = await createAndRenderPreview();
@@ -535,7 +535,7 @@ describe('PreviewWeb', () => {
 
       it('renders exception if the story calls showException', async () => {
         const error = new Error('error');
-        projectAnnotations.renderToDOM.mockImplementationOnce((context) =>
+        projectAnnotations.renderToDOM.mockImplementation((context) =>
           context.showException(error)
         );
 
@@ -562,7 +562,7 @@ describe('PreviewWeb', () => {
 
       it('does not show error display if the render function throws IGNORED_EXCEPTION', async () => {
         document.location.search = '?id=component-one--a';
-        projectAnnotations.renderToDOM.mockImplementationOnce(() => {
+        projectAnnotations.renderToDOM.mockImplementation(() => {
           throw IGNORED_EXCEPTION;
         });
 
@@ -837,7 +837,7 @@ describe('PreviewWeb', () => {
         const [gate, openGate] = createGate();
 
         document.location.search = '?id=component-one--a';
-        projectAnnotations.renderToDOM.mockImplementationOnce(async () => gate);
+        projectAnnotations.renderToDOM.mockImplementation(async () => gate);
 
         await new PreviewWeb().initialize({ importFn, getProjectAnnotations });
         await waitForRenderPhase('rendering');
@@ -876,7 +876,7 @@ describe('PreviewWeb', () => {
 
       it('works if it is called directly from inside non async renderToDOM', async () => {
         document.location.search = '?id=component-one--a';
-        projectAnnotations.renderToDOM.mockImplementationOnce(() => {
+        projectAnnotations.renderToDOM.mockImplementation(() => {
           emitter.emit(UPDATE_STORY_ARGS, {
             storyId: 'component-one--a',
             updatedArgs: { new: 'arg' },
@@ -912,7 +912,7 @@ describe('PreviewWeb', () => {
         componentOneExports.a.play.mockImplementationOnce(async () => gate);
 
         const renderToDOMCalled = new Promise((resolve) => {
-          projectAnnotations.renderToDOM.mockImplementationOnce(() => {
+          projectAnnotations.renderToDOM.mockImplementation(() => {
             resolve(null);
           });
         });
@@ -1296,7 +1296,7 @@ describe('PreviewWeb', () => {
       const [gate, openGate] = createGate();
 
       document.location.search = '?id=component-one--a';
-      projectAnnotations.renderToDOM.mockImplementationOnce(async () => gate);
+      projectAnnotations.renderToDOM.mockImplementation(async () => gate);
       await new PreviewWeb().initialize({ importFn, getProjectAnnotations });
       await waitForRenderPhase('rendering');
 
@@ -1463,6 +1463,20 @@ describe('PreviewWeb', () => {
         expect(projectAnnotations.renderToDOM).not.toHaveBeenCalled();
       });
 
+      it('does NOT call renderToDOMs teardown', async () => {
+        document.location.search = '?id=component-one--a';
+        await createAndRenderPreview();
+
+        projectAnnotations.renderToDOM.mockClear();
+        emitter.emit(SET_CURRENT_STORY, {
+          storyId: 'component-one--a',
+          viewMode: 'story',
+        });
+        await waitForSetCurrentStory();
+
+        expect(teardownRenderToDOM).not.toHaveBeenCalled();
+      });
+
       // For https://github.com/storybookjs/storybook/issues/17214
       it('does NOT render a second time if preparing', async () => {
         document.location.search = '?id=component-one--a';
@@ -1510,6 +1524,20 @@ describe('PreviewWeb', () => {
     });
 
     describe('when changing story in story viewMode', () => {
+      it('calls renderToDOMs teardown', async () => {
+        document.location.search = '?id=component-one--a';
+        await createAndRenderPreview();
+
+        projectAnnotations.renderToDOM.mockClear();
+        emitter.emit(SET_CURRENT_STORY, {
+          storyId: 'component-one--b',
+          viewMode: 'story',
+        });
+        await waitForSetCurrentStory();
+
+        expect(teardownRenderToDOM).toHaveBeenCalled();
+      });
+
       it('updates URL', async () => {
         document.location.search = '?id=component-one--a';
         await createAndRenderPreview();
@@ -1645,7 +1673,7 @@ describe('PreviewWeb', () => {
         const preview = await createAndRenderPreview();
 
         const error = new Error('error');
-        projectAnnotations.renderToDOM.mockImplementationOnce(() => {
+        projectAnnotations.renderToDOM.mockImplementation(() => {
           throw error;
         });
 
@@ -1666,9 +1694,7 @@ describe('PreviewWeb', () => {
         const preview = await createAndRenderPreview();
 
         const error = { title: 'title', description: 'description' };
-        projectAnnotations.renderToDOM.mockImplementationOnce((context) =>
-          context.showError(error)
-        );
+        projectAnnotations.renderToDOM.mockImplementation((context) => context.showError(error));
 
         mockChannel.emit.mockClear();
         emitter.emit(SET_CURRENT_STORY, {
@@ -1690,7 +1716,7 @@ describe('PreviewWeb', () => {
         const preview = await createAndRenderPreview();
 
         const error = new Error('error');
-        projectAnnotations.renderToDOM.mockImplementationOnce((context) =>
+        projectAnnotations.renderToDOM.mockImplementation((context) =>
           context.showException(error)
         );
 
@@ -1811,7 +1837,7 @@ describe('PreviewWeb', () => {
           const [gate, openGate] = createGate();
 
           document.location.search = '?id=component-one--a';
-          projectAnnotations.renderToDOM.mockImplementationOnce(async () => gate);
+          projectAnnotations.renderToDOM.mockImplementation(async () => gate);
           await new PreviewWeb().initialize({ importFn, getProjectAnnotations });
           await waitForRenderPhase('rendering');
 
@@ -1937,6 +1963,19 @@ describe('PreviewWeb', () => {
     });
 
     describe('when changing from story viewMode to docs', () => {
+      it('calls renderToDOMs teardown', async () => {
+        document.location.search = '?id=component-one--a';
+        await createAndRenderPreview();
+
+        emitter.emit(SET_CURRENT_STORY, {
+          storyId: 'component-one--a',
+          viewMode: 'docs',
+        });
+        await waitForSetCurrentStory();
+
+        expect(teardownRenderToDOM).toHaveBeenCalled();
+      });
+
       it('updates URL', async () => {
         document.location.search = '?id=component-one--a';
         await createAndRenderPreview();
@@ -2199,7 +2238,7 @@ describe('PreviewWeb', () => {
         const preview = await createAndRenderPreview();
 
         const error = new Error('error');
-        projectAnnotations.renderToDOM.mockImplementationOnce(() => {
+        projectAnnotations.renderToDOM.mockImplementation(() => {
           throw error;
         });
 
@@ -2217,9 +2256,7 @@ describe('PreviewWeb', () => {
 
       it('renders error if the story calls showError', async () => {
         const error = { title: 'title', description: 'description' };
-        projectAnnotations.renderToDOM.mockImplementationOnce((context) =>
-          context.showError(error)
-        );
+        projectAnnotations.renderToDOM.mockImplementation((context) => context.showError(error));
 
         document.location.search = '?id=component-one--a&viewMode=docs';
         const preview = await createAndRenderPreview();
@@ -2241,7 +2278,7 @@ describe('PreviewWeb', () => {
 
       it('renders exception if the story calls showException', async () => {
         const error = new Error('error');
-        projectAnnotations.renderToDOM.mockImplementationOnce((context) =>
+        projectAnnotations.renderToDOM.mockImplementation((context) =>
           context.showException(error)
         );
 
@@ -2347,6 +2384,17 @@ describe('PreviewWeb', () => {
           : componentTwoExports;
       });
 
+      it('calls renderToDOMs teardown', async () => {
+        document.location.search = '?id=component-one--a';
+        const preview = await createAndRenderPreview();
+        mockChannel.emit.mockClear();
+
+        preview.onStoriesChanged({ importFn: newImportFn });
+        await waitForRender();
+
+        expect(teardownRenderToDOM).toHaveBeenCalled();
+      });
+
       it('does not emit STORY_UNCHANGED', async () => {
         document.location.search = '?id=component-one--a';
         const preview = await createAndRenderPreview();
@@ -2491,7 +2539,7 @@ describe('PreviewWeb', () => {
         const preview = await createAndRenderPreview();
 
         const error = new Error('error');
-        projectAnnotations.renderToDOM.mockImplementationOnce(() => {
+        projectAnnotations.renderToDOM.mockImplementation(() => {
           throw error;
         });
 
@@ -2508,9 +2556,7 @@ describe('PreviewWeb', () => {
         const preview = await createAndRenderPreview();
 
         const error = { title: 'title', description: 'description' };
-        projectAnnotations.renderToDOM.mockImplementationOnce((context) =>
-          context.showError(error)
-        );
+        projectAnnotations.renderToDOM.mockImplementation((context) => context.showError(error));
 
         mockChannel.emit.mockClear();
         preview.onStoriesChanged({ importFn: newImportFn });
@@ -2528,7 +2574,7 @@ describe('PreviewWeb', () => {
         const preview = await createAndRenderPreview();
 
         const error = new Error('error');
-        projectAnnotations.renderToDOM.mockImplementationOnce((context) =>
+        projectAnnotations.renderToDOM.mockImplementation((context) =>
           context.showException(error)
         );
 
@@ -2629,6 +2675,17 @@ describe('PreviewWeb', () => {
           : newComponentTwoExports;
       });
 
+      it('does NOT call renderToDOMs teardown', async () => {
+        document.location.search = '?id=component-one--a';
+        const preview = await createAndRenderPreview();
+
+        mockChannel.emit.mockClear();
+        preview.onStoriesChanged({ importFn: newImportFn });
+        await waitForEvents([STORY_UNCHANGED]);
+
+        expect(teardownRenderToDOM).not.toHaveBeenCalled();
+      });
+
       it('emits STORY_UNCHANGED', async () => {
         document.location.search = '?id=component-one--a';
         const preview = await createAndRenderPreview();
@@ -2754,6 +2811,17 @@ describe('PreviewWeb', () => {
         },
       };
 
+      it('calls renderToDOMs teardown', async () => {
+        document.location.search = '?id=component-one--a';
+        const preview = await createAndRenderPreview();
+
+        mockChannel.emit.mockClear();
+        preview.onStoriesChanged({ importFn: newImportFn, storyIndex: newStoryIndex });
+        await waitForEvents([STORY_MISSING]);
+
+        expect(teardownRenderToDOM).toHaveBeenCalled();
+      });
+
       it('renders loading error', async () => {
         document.location.search = '?id=component-one--a';
         const preview = await createAndRenderPreview();
@@ -2920,6 +2988,18 @@ describe('PreviewWeb', () => {
       });
     });
 
+    it('calls renderToDOMs teardown', async () => {
+      document.location.search = '?id=component-one--a';
+      const preview = await createAndRenderPreview();
+
+      projectAnnotations.renderToDOM.mockClear();
+      mockChannel.emit.mockClear();
+      preview.onGetProjectAnnotationsChanged({ getProjectAnnotations: newGetProjectAnnotations });
+      await waitForRender();
+
+      expect(teardownRenderToDOM).toHaveBeenCalled();
+    });
+
     it('rerenders the current story with new global meta-generated context', async () => {
       document.location.search = '?id=component-one--a';
       const preview = await createAndRenderPreview();

lib/preview-web/src/StoryRender.ts

@@ -6,7 +6,13 @@ import {
   StoryContextForLoaders,
   StoryContext,
 } from '@storybook/csf';
-import { Story, RenderContext, StoryStore } from '@storybook/store';
+import {
+  Story,
+  RenderContext,
+  StoryStore,
+  RenderToDOM,
+  TeardownRenderToDOM,
+} from '@storybook/store';
 import { Channel } from '@storybook/addons';
 import { STORY_RENDER_PHASE_CHANGED, STORY_RENDERED } from '@storybook/core-events';
 
@@ -62,13 +68,12 @@ export class StoryRender<TFramework extends AnyFramework> implements Render<TFra
 
   public disableKeyListeners = false;
 
+  private teardownRender: TeardownRenderToDOM = () => {};
+
   constructor(
     public channel: Channel,
     public store: StoryStore<TFramework>,
-    private renderToScreen: (
-      renderContext: RenderContext<TFramework>,
-      canvasElement: HTMLElement
-    ) => void | Promise<void>,
+    private renderToScreen: RenderToDOM<TFramework>,
     private callbacks: RenderContextCallbacks<TFramework>,
     public id: StoryId,
     public viewMode: ViewMode,
@@ -190,9 +195,10 @@ export class StoryRender<TFramework extends AnyFramework> implements Render<TFra
         unboundStoryFn,
       };
 
-      await this.runPhase(abortSignal, 'rendering', async () =>
-        this.renderToScreen(renderContext, this.canvasElement)
-      );
+      await this.runPhase(abortSignal, 'rendering', async () => {
+        this.teardownRender =
+          (await this.renderToScreen(renderContext, this.canvasElement)) || (() => {});
+      });
       this.notYetRendered = false;
       if (abortSignal.aborted) return;
 
@@ -241,7 +247,11 @@ export class StoryRender<TFramework extends AnyFramework> implements Render<TFra
     // Wait several ticks that may be needed to handle the abort, then try again.
     // Note that there's a max of 5 nested timeouts before they're no longer "instant".
     for (let i = 0; i < 3; i += 1) {
-      if (!this.isPending()) return;
+      if (!this.isPending()) {
+        // eslint-disable-next-line no-await-in-loop
+        await this.teardownRender();
+        return;
+      }
       // eslint-disable-next-line no-await-in-loop
       await new Promise((resolve) => setTimeout(resolve, 0));
     }

lib/store/src/types.ts

@@ -29,9 +29,17 @@ export type ModuleExports = Record<string, any>;
 export type PromiseLike<T> = Promise<T> | SynchronousPromise<T>;
 export type ModuleImportFn = (path: Path) => PromiseLike<ModuleExports>;
 
+type MaybePromise<T> = Promise<T> | T;
+
+export type TeardownRenderToDOM = () => MaybePromise<void>;
+export type RenderToDOM<TFramework extends AnyFramework> = (
+  context: RenderContext<TFramework>,
+  element: Element
+) => MaybePromise<void | TeardownRenderToDOM>;
+
 export type WebProjectAnnotations<TFramework extends AnyFramework> =
   ProjectAnnotations<TFramework> & {
-    renderToDOM?: (context: RenderContext<TFramework>, element: Element) => Promise<void> | void;
+    renderToDOM?: RenderToDOM<TFramework>;
   };
 
 export type NormalizedProjectAnnotations<TFramework extends AnyFramework = AnyFramework> =

netlify.toml

@@ -2,7 +2,7 @@
   publish = "built-storybooks"
   command = "yarn bootstrap --core && yarn build-storybooks --all"
 [build.environment]
-  NODE_VERSION = "12"
+  NODE_VERSION = "14"
   YARN_VERSION = "1.22.10"
   DOTENV_DISPLAY_WARNING = "none"
   STORYBOOK_EXAMPLE_APP ="true"