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/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/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> =