From 8e2a37594fb564112754e36cef20645ca82960b8 Mon Sep 17 00:00:00 2001
From: Yann Braga <yannbf@gmail.com>
Date: Thu, 8 Feb 2024 13:05:51 +0100
Subject: [PATCH] pass story context to the play function of a composed story

---
 .../store/csf/portable-stories.test.ts        | 32 ++++++++++++++++++
 .../src/modules/store/csf/portable-stories.ts | 33 +++++++++++++++----
 2 files changed, 58 insertions(+), 7 deletions(-)

diff --git a/code/lib/preview-api/src/modules/store/csf/portable-stories.test.ts b/code/lib/preview-api/src/modules/store/csf/portable-stories.test.ts
index 0531a5acdc18..eae92f09aac8 100644
--- a/code/lib/preview-api/src/modules/store/csf/portable-stories.test.ts
+++ b/code/lib/preview-api/src/modules/store/csf/portable-stories.test.ts
@@ -30,6 +30,38 @@ describe('composeStory', () => {
     );
   });
 
+  it('should compose with a play function', async () => {
+    const spy = vi.fn();
+    const Story = () => {};
+    Story.args = {
+      primary: true,
+    };
+    Story.play = async (context: any) => {
+      spy(context);
+    };
+
+    const composedStory = composeStory(Story, meta);
+    await composedStory.play({ canvasElement: null });
+    expect(spy).toHaveBeenCalledWith(
+      expect.objectContaining({
+        args: {
+          ...Story.args,
+          ...meta.args,
+        },
+      })
+    );
+  });
+
+  it('should throw when executing the play function but the story does not have one', async () => {
+    const Story = () => {};
+    Story.args = {
+      primary: true,
+    };
+
+    const composedStory = composeStory(Story, meta);
+    expect(composedStory.play({ canvasElement: null })).rejects.toThrow();
+  });
+
   it('should throw an error if Story is undefined', () => {
     expect(() => {
       // @ts-expect-error (invalid input)
diff --git a/code/lib/preview-api/src/modules/store/csf/portable-stories.ts b/code/lib/preview-api/src/modules/store/csf/portable-stories.ts
index b9d7abc31213..0ebf858e96e8 100644
--- a/code/lib/preview-api/src/modules/store/csf/portable-stories.ts
+++ b/code/lib/preview-api/src/modules/store/csf/portable-stories.ts
@@ -12,6 +12,7 @@ import type {
   Parameters,
   ComposedStoryFn,
   StrictArgTypes,
+  ComposedStoryPlayContext,
 } from '@storybook/types';
 
 import { HooksContext } from '../../../addons';
@@ -74,24 +75,42 @@ export function composeStory<TRenderer extends Renderer = Renderer, TArgs extend
 
   const defaultGlobals = getValuesFromArgTypes(projectAnnotations.globalTypes);
 
+  const context: StoryContext<TRenderer> = {
+    hooks: new HooksContext(),
+    globals: defaultGlobals,
+    args: { ...story.initialArgs },
+    viewMode: 'story',
+    loaded: {},
+    abortSignal: null as unknown as AbortSignal,
+    canvasElement: null,
+    ...story,
+  };
+
   const composedStory: ComposedStoryFn<TRenderer, Partial<TArgs>> = Object.assign(
     (extraArgs?: Partial<TArgs>) => {
-      const context: Partial<StoryContext> = {
-        ...story,
-        hooks: new HooksContext(),
-        globals: defaultGlobals,
-        args: { ...story.initialArgs, ...extraArgs },
+      const finalContext: StoryContext<TRenderer> = {
+        ...context,
+        args: { ...context.initialArgs, ...extraArgs },
       };
 
-      return story.unboundStoryFn(prepareContext(context as StoryContext<TRenderer>));
+      return story.unboundStoryFn(prepareContext(finalContext));
     },
     {
       storyName,
       args: story.initialArgs as Partial<TArgs>,
-      play: story.playFunction as ComposedStoryPlayFn<TRenderer, Partial<TArgs>>,
       parameters: story.parameters as Parameters,
       argTypes: story.argTypes as StrictArgTypes<TArgs>,
       id: story.id,
+      play: (async (extraContext: ComposedStoryPlayContext<TRenderer, TArgs>) => {
+        if (story.playFunction === undefined) {
+          throw new Error('The story does not have a play function. Make sure to add one.');
+        }
+
+        await story.playFunction({
+          ...context,
+          ...extraContext,
+        });
+      }) as unknown as ComposedStoryPlayFn<TRenderer, Partial<TArgs>>,
     }
   );