Pass playContext to step function

storybookjs · Jul 8, 2022 · 640072e · 640072e
1 parent 63d2f98
commit 640072e
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 14 deletions.
diff --git a/lib/preview-web/src/StoryRender.ts b/lib/preview-web/src/StoryRender.ts
@@ -9,6 +9,7 @@ import {
 import {
   Story,
   RenderContext,
+  PlayContext,
   StoryStore,
   RenderToDOM,
   TeardownRenderToDOM,
@@ -48,6 +49,20 @@ function createController(): AbortController {
   } as AbortController;
 }
 
+function createStepFunction(context: PlayContext) {
+  const { step } = instrument(
+    { step: (label: string, play: (context: PlayContext) => MaybePromise<void>) => play(context) },
+    { intercept: true }
+  );
+  return step;
+}
+
+function createPlayContext(storyContext: StoryContext): PlayContext {
+  const playContext = { ...storyContext } as any;
+  playContext.step = createStepFunction(playContext);
+  return playContext;
+}
+
 function serializeError(error: any) {
   try {
     const { name = 'Error', message = String(error), stack } = error;
@@ -57,14 +72,6 @@ function serializeError(error: any) {
   }
 }
 
-function createStepFunction() {
-  const { step } = instrument(
-    { step: (label: string, callback: () => MaybePromise<void>) => callback() },
-    { intercept: true }
-  );
-  return step;
-}
-
 export type RenderContextCallbacks<TFramework extends AnyFramework> = Pick<
   RenderContext<TFramework>,
   'showMain' | 'showError' | 'showException'
@@ -93,7 +100,7 @@ export class StoryRender<TFramework extends AnyFramework> implements Render<TFra
 
   private abortController?: AbortController;
 
-  private stepFunction?: (label: string, callback: () => MaybePromise<void>) => void;
+  private stepFunction?: PlayContext['step'];
 
   private canvasElement?: HTMLElement;
 
@@ -115,7 +122,6 @@ export class StoryRender<TFramework extends AnyFramework> implements Render<TFra
     story?: Story<TFramework>
   ) {
     this.abortController = createController();
-    this.stepFunction = createStepFunction();
 
     // Allow short-circuiting preparing if we happen to already
     // have the story (this is used by docs mode)
@@ -240,7 +246,7 @@ export class StoryRender<TFramework extends AnyFramework> implements Render<TFra
           return this.callbacks.showException(error);
         },
         forceRemount: forceRemount || this.notYetRendered,
-        playContext: { ...renderStoryContext, step: this.stepFunction },
+        playContext: createPlayContext(renderStoryContext),
         storyContext: renderStoryContext,
         storyFn: () => unboundStoryFn(renderStoryContext),
         unboundStoryFn,

diff --git a/lib/store/src/types.ts b/lib/store/src/types.ts
@@ -106,15 +106,21 @@ export type BoundStory<TFramework extends AnyFramework = AnyFramework> = Story<T
   storyFn: PartialStoryFn<TFramework>;
 };
 
+export declare type PlayContext<TFramework extends AnyFramework = AnyFramework> =
+  StoryContext<TFramework> & {
+    step: (
+      name: string,
+      play: (context: PlayContext<TFramework>) => MaybePromise<void>
+    ) => MaybePromise<void>;
+  };
+
 export declare type RenderContext<TFramework extends AnyFramework = AnyFramework> =
   StoryIdentifier & {
     showMain: () => void;
     showError: (error: { title: string; description: string }) => void;
     showException: (err: Error) => void;
     forceRemount: boolean;
-    playContext: StoryContext<TFramework> & {
-      step: (label: string, callback: () => MaybePromise<void>) => MaybePromise<void>;
-    };
+    playContext: PlayContext<TFramework>;
     storyContext: StoryContext<TFramework>;
     storyFn: PartialStoryFn<TFramework>;
     unboundStoryFn: LegacyStoryFn<TFramework>;