Improvements

storybookjs · Jul 11, 2022 · a1b36f5 · a1b36f5
1 parent 8acb04b
commit a1b36f5
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 14 deletions.
diff --git a/lib/store/src/csf/prepareStory.ts b/lib/store/src/csf/prepareStory.ts
@@ -13,6 +13,8 @@ import type {
   StrictArgTypes,
   StoryContextForLoaders,
   PlayFunctionContext,
+  StepLabel,
+  PlayFunction,
 } from '@storybook/csf';
 import { includeConditionalArg } from '@storybook/csf';
 
@@ -205,8 +207,9 @@ export function prepareStory<TFramework extends AnyFramework>(
     (async (storyContext: StoryContext<TFramework>) => {
       const playFunctionContext: PlayFunctionContext<TFramework> = {
         ...storyContext,
-        // TODO: We know runStep is defined, we need a proper normalized annotations type
-        step: (label, play) => runStep!(label, play, playFunctionContext),
+        step: (label: StepLabel, play: PlayFunction<TFramework>) =>
+          // TODO: We know runStep is defined, we need a proper normalized annotations type
+          runStep!(label, play, playFunctionContext),
       };
       return play(playFunctionContext);
     });

diff --git a/lib/store/src/csf/stepRunners.test.ts b/lib/store/src/csf/stepRunners.test.ts
@@ -3,16 +3,18 @@ import { composeStepRunners } from './stepRunners';
 
 describe('stepRunners', () => {
   it('composes each step runner', async () => {
-    const firstLabels: string[] = [];
+    const order: string[] = [];
+
     const firstStepRunner: StepRunner = async (label, play, ctx) => {
-      firstLabels.push(label);
-      return play(ctx);
+      order.push(`first-${label}-start`);
+      await play(ctx);
+      order.push(`first-${label}-end`);
     };
 
-    const secondLabels: string[] = [];
     const secondStepRunner: StepRunner = async (label, play, ctx) => {
-      secondLabels.push(label);
-      return play(ctx);
+      order.push(`second-${label}-start`);
+      await play(ctx);
+      order.push(`second-${label}-end`);
     };
 
     const composed = composeStepRunners([firstStepRunner, secondStepRunner]);
@@ -28,8 +30,16 @@ describe('stepRunners', () => {
     expect(playFnA).toHaveBeenCalledWith(playContextA);
     expect(playFnB).toHaveBeenCalledTimes(1);
     expect(playFnB).toHaveBeenCalledWith(playContextB);
-    expect(firstLabels).toEqual(['a', 'b']);
-    expect(secondLabels).toEqual(['a', 'b']);
+    expect(order).toEqual([
+      'first-a-start',
+      'second-a-start',
+      'second-a-end',
+      'first-a-end',
+      'first-b-start',
+      'second-b-start',
+      'second-b-end',
+      'first-b-end',
+    ]);
   });
 
   it('creates a sensible default if no step runner is provided', async () => {

diff --git a/lib/store/src/csf/stepRunners.ts b/lib/store/src/csf/stepRunners.ts
@@ -22,10 +22,10 @@ export function composeStepRunners<TFramework extends AnyFramework>(
   stepRunners: StepRunner<TFramework>[]
 ): StepRunner<TFramework> {
   return async (label, play, playContext) => {
-    const composedPlay = await stepRunners.reduce(
-      async (innerPlay, stepRunner) => async () => stepRunner(label, await innerPlay, playContext),
-      Promise.resolve(play)
+    const composedPlay = stepRunners.reduceRight<() => Promise<void>>(
+      (innerPlay, stepRunner) => async () => stepRunner(label, innerPlay, playContext),
+      async () => play(playContext)
     );
-    await composedPlay(playContext);
+    await composedPlay();
   };
 }