Core: Refactor phases to run in order loading -> rendering -> playing

code/core/src/preview-api/modules/preview-web/PreviewWithSelection.tsx

@@ -308,11 +308,7 @@ export class PreviewWithSelection<TRenderer extends Renderer> extends Preview<TR
       render = new StoryRender<TRenderer>(
         this.channel,
         this.storyStoreValue,
-        (...args: Parameters<typeof renderToCanvas>) => {
-          // At the start of renderToCanvas we make the story visible (see note in WebView)
-          this.view.showStoryDuringRender();
-          return renderToCanvas(...args);
-        },
+        renderToCanvas,
         this.mainStoryCallbacks(storyId),
         storyId,
         'story'
@@ -434,6 +430,7 @@ export class PreviewWithSelection<TRenderer extends Renderer> extends Preview<TR
   // UTILITIES
   mainStoryCallbacks(storyId: StoryId) {
     return {
+      showStoryDuringRender: () => this.view.showStoryDuringRender(),
       showMain: () => this.view.showMain(),
       showError: (err: { title: string; description: string }) => this.renderError(storyId, err),
       showException: (err: Error) => this.renderException(storyId, err),

code/core/src/preview-api/modules/preview-web/render/StoryRender.test.ts

@@ -1,5 +1,5 @@
 // @vitest-environment happy-dom
-import { describe, it, expect, vi } from 'vitest';
+import { describe, it, expect, vi, beforeEach } from 'vitest';
 import { Channel } from '@storybook/core/channels';
 import type { PreparedStory, Renderer, StoryContext, StoryIndexEntry } from '@storybook/core/types';
 import type { StoryStore } from '../../store';
@@ -26,6 +26,11 @@ const tick = () => new Promise((resolve) => setTimeout(resolve, 0));
 
 window.location = { reload: vi.fn() } as any;
 
+const mountSpy = vi.fn(async (context) => {
+  await context.renderToCanvas();
+  return context.canvas;
+});
+
 const buildStory = (overrides: Partial<PreparedStory> = {}): PreparedStory =>
   ({
     id: 'id',
@@ -36,12 +41,7 @@ const buildStory = (overrides: Partial<PreparedStory> = {}): PreparedStory =>
     applyBeforeEach: vi.fn(),
     unboundStoryFn: vi.fn(),
     playFunction: vi.fn(),
-    mount: vi.fn((context: StoryContext) => {
-      return async () => {
-        await context.renderToCanvas();
-        return context.canvas;
-      };
-    }),
+    mount: (context: StoryContext) => () => mountSpy(context),
     ...overrides,
   }) as any;
 
@@ -53,6 +53,10 @@ const buildStore = (overrides: Partial<StoryStore<Renderer>> = {}): StoryStore<R
     ...overrides,
   }) as any;
 
+beforeEach(() => {
+  vi.restoreAllMocks();
+});
+
 describe('StoryRender', () => {
   it('does run play function if passed autoplay=true', async () => {
     const story = buildStory();
@@ -133,12 +137,7 @@ describe('StoryRender', () => {
   });
 
   it('calls mount if play function does not destructure mount', async () => {
-    const actualMount = vi.fn(async (context) => {
-      await context.renderToCanvas();
-      return {};
-    });
     const story = buildStory({
-      mount: (context) => () => actualMount(context) as any,
       playFunction: () => {},
     });
     const render = new StoryRender(
@@ -153,17 +152,14 @@ describe('StoryRender', () => {
     );
 
     await render.renderToElement({} as any);
-    expect(actualMount).toHaveBeenCalled();
+    expect(mountSpy).toHaveBeenCalledOnce();
   });
 
-  it('does not call mount if play function destructures mount', async () => {
-    const actualMount = vi.fn(async (context) => {
-      await context.renderToCanvas();
-      return context.canvas;
-    });
+  it('does not call mount twice if mount called in play function', async () => {
     const story = buildStory({
-      mount: (context) => () => actualMount(context) as any,
-      playFunction: ({ mount }) => {},
+      playFunction: async ({ mount }) => {
+        await mount();
+      },
     });
     const render = new StoryRender(
       new Channel({}),
@@ -177,16 +173,11 @@ describe('StoryRender', () => {
     );
 
     await render.renderToElement({} as any);
-    expect(actualMount).not.toHaveBeenCalled();
+    expect(mountSpy).toHaveBeenCalledOnce();
   });
 
   it('errors if play function calls mount without destructuring', async () => {
-    const actualMount = vi.fn(async (context) => {
-      await context.renderToCanvas();
-      return {};
-    });
     const story = buildStory({
-      mount: (context) => () => actualMount(context) as any,
       playFunction: async (context) => {
         await context.mount();
       },
@@ -207,16 +198,40 @@ describe('StoryRender', () => {
     expect(view.showException).toHaveBeenCalled();
   });
 
+  it('errors if play function destructures mount but does not call it', async () => {
+    const story = buildStory({
+      playFunction: async ({ mount }) => {
+        // forget to call mount
+      },
+    });
+    const view = { showException: vi.fn() };
+    const render = new StoryRender(
+      new Channel({}),
+      buildStore(),
+      vi.fn() as any,
+      view as any,
+      entry.id,
+      'story',
+      { autoplay: true },
+      story
+    );
+
+    await render.renderToElement({} as any);
+    expect(view.showException).toHaveBeenCalled();
+  });
+
   it('enters rendering phase during play if play function calls mount', async () => {
     const actualMount = vi.fn(async (context) => {
       await context.renderToCanvas();
-      return {};
+      expect(render.phase).toBe('rendering');
+      return context.canvas;
     });
     const story = buildStory({
       mount: (context) => () => actualMount(context) as any,
-      playFunction: ({ mount }) => {
+      playFunction: async ({ mount }) => {
+        expect(render.phase).toBe('loading');
+        await mount();
         expect(render.phase).toBe('playing');
-        mount();
       },
     });
     const render = new StoryRender(

code/core/src/preview-api/modules/preview-web/render/StoryRender.ts

@@ -9,7 +9,7 @@ import type { StoryStore } from '../../store';
 import type { Render, RenderType } from './Render';
 import { PREPARE_ABORTED } from './Render';
 import { mountDestructured } from './mount-utils';
-import { MountMustBeDestructuredError } from '@storybook/core-events/preview-errors';
+import { MountMustBeDestructuredError, NoStoryMountedError } from '@storybook/core/preview-errors';
 
 import type {
   Canvas,
@@ -71,7 +71,7 @@ export class StoryRender<TRenderer extends Renderer> implements Render<TRenderer
     public channel: Channel,
     public store: StoryStore<TRenderer>,
     private renderToScreen: RenderToCanvas<TRenderer>,
-    private callbacks: RenderContextCallbacks<TRenderer>,
+    private callbacks: RenderContextCallbacks<TRenderer> & { showStoryDuringRender?: () => void },
     public id: StoryId,
     public viewMode: StoryContext['viewMode'],
     public renderOptions: StoryRenderOptions = { autoplay: true, forceInitialArgs: false },
@@ -92,12 +92,19 @@ export class StoryRender<TRenderer extends Renderer> implements Render<TRenderer
   private async runPhase(signal: AbortSignal, phase: RenderPhase, phaseFn?: () => Promise<void>) {
     this.phase = phase;
     this.channel.emit(STORY_RENDER_PHASE_CHANGED, { newPhase: this.phase, storyId: this.id });
-    if (phaseFn) await phaseFn();
+    if (phaseFn) {
+      await phaseFn();
+      this.checkIfAborted(signal);
+    }
+  }
 
+  private checkIfAborted(signal: AbortSignal): boolean {
     if (signal.aborted) {
       this.phase = 'aborted';
       this.channel.emit(STORY_RENDER_PHASE_CHANGED, { newPhase: this.phase, storyId: this.id });
+      return true;
     }
+    return false;
   }
 
   async prepare() {
@@ -196,22 +203,29 @@ export class StoryRender<TRenderer extends Renderer> implements Render<TRenderer
         step: (label, play) => runStep(label, play, context),
         context: null!,
         canvas: {} as Canvas,
-        mount: null!,
         renderToCanvas: async () => {
+          const teardown = await this.renderToScreen(renderContext, canvasElement);
+          this.teardownRender = teardown || (() => {});
+          mounted = true;
+        },
+        // The story provides (set in a renderer) a mount function that is a higher order function
+        // (context) => (...args) => Canvas
+        //
+        // Before assigning it to the context, we resolve the context dependency,
+        // so that a user can just call it as await mount(...args) in their play function.
+        mount: async (...args) => {
+          this.callbacks.showStoryDuringRender?.();
+          let mountReturn: Awaited<ReturnType<StoryContext['mount']>> = null!;
           await this.runPhase(abortSignal, 'rendering', async () => {
-            const teardown = await this.renderToScreen(renderContext, canvasElement);
-            this.teardownRender = teardown || (() => {});
-            mounted = true;
+            mountReturn = await story.mount(context)(...args);
           });
-
-          if (isMountDestructured) {
-            // put the phase back to playing if mount is used inside a play function
-            await this.runPhase(abortSignal, 'playing', async () => {});
-          }
+          // start playing phase if mount is used inside a play function
+          if (isMountDestructured) await this.runPhase(abortSignal, 'playing');
+          return mountReturn;
         },
       };
+
       context.context = context;
-      context.mount = this.story.mount(context);
 
       const renderContext: RenderContext<TRenderer> = {
         componentId,
@@ -241,12 +255,10 @@ export class StoryRender<TRenderer extends Renderer> implements Render<TRenderer
 
       if (abortSignal.aborted) return;
 
-      await this.runPhase(abortSignal, 'beforeEach', async () => {
-        const cleanupCallbacks = await applyBeforeEach(context);
-        this.store.addCleanupCallbacks(story, cleanupCallbacks);
-      });
+      const cleanupCallbacks = await applyBeforeEach(context);
+      this.store.addCleanupCallbacks(story, cleanupCallbacks);
 
-      if (abortSignal.aborted) return;
+      if (this.checkIfAborted(abortSignal)) return;
 
       if (!mounted && !isMountDestructured) {
         await context.mount();
@@ -272,16 +284,26 @@ export class StoryRender<TRenderer extends Renderer> implements Render<TRenderer
             context.mount = async () => {
               throw new MountMustBeDestructuredError({ playFunction: playFunction.toString() });
             };
-          }
-          await this.runPhase(abortSignal, 'playing', async () => {
+            await this.runPhase(abortSignal, 'playing', async () => playFunction(context));
+          } else {
+            // when mount is used the playing phase will start later, right after mount is called in the play function
             await playFunction(context);
-          });
+          }
+
+          if (!mounted) {
+            throw new NoStoryMountedError();
+          }
+          this.checkIfAborted(abortSignal);
+
           if (!ignoreUnhandledErrors && unhandledErrors.size > 0) {
             await this.runPhase(abortSignal, 'errored');
           } else {
             await this.runPhase(abortSignal, 'played');
           }
         } catch (error) {
+          // Remove the loading screen, even if there was an error before rendering
+          this.callbacks.showStoryDuringRender?.();
+
           await this.runPhase(abortSignal, 'errored', async () => {
             this.channel.emit(PLAY_FUNCTION_THREW_EXCEPTION, serializeError(error));
           });

code/core/src/preview-errors.ts

@@ -295,6 +295,29 @@ export class TestingLibraryMustBeConfiguredError extends StorybookError {
   }
 }
 
+export class NoStoryMountedError extends StorybookError {
+  readonly category = Category.PREVIEW_API;
+
+  readonly code = 14;
+
+  template() {
+    return dedent`
+    No story is mounted in your story.
+
+    This usually occurs when you destructure mount in the play function, but forgot to call it.
+
+    For example:
+
+    async play({ mount, canvasElement }) {
+      // 👈mount should be called: await mount(); 
+      const canvas = within(canvasElement);
+      const button = await canvas.findByRole('button');
+      await userEvent.click(button);
+    };
+    `;
+  }
+}
+
 export class NextJsSharpError extends StorybookError {
   readonly category = Category.FRAMEWORK_NEXTJS;