From 0d176deb9a7e61ae8c2c5d25582d3e4049fbe048 Mon Sep 17 00:00:00 2001
From: Tom Coleman <tom@chromatic.com>
Date: Mon, 11 Jul 2022 11:45:13 +1000
Subject: [PATCH] Improvements

---
 code/lib/store/src/csf/prepareStory.ts     |  7 ++++--
 code/lib/store/src/csf/stepRunners.test.ts | 26 +++++++++++++++-------
 code/lib/store/src/csf/stepRunners.ts      |  8 +++----
 3 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/code/lib/store/src/csf/prepareStory.ts b/code/lib/store/src/csf/prepareStory.ts
index 115623488c22..92f8ce6e3e71 100644
--- a/code/lib/store/src/csf/prepareStory.ts
+++ b/code/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/code/lib/store/src/csf/stepRunners.test.ts b/code/lib/store/src/csf/stepRunners.test.ts
index ad57f47a2e67..5a65b59acab6 100644
--- a/code/lib/store/src/csf/stepRunners.test.ts
+++ b/code/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/code/lib/store/src/csf/stepRunners.ts b/code/lib/store/src/csf/stepRunners.ts
index 318b9270ea75..af12acd12d4d 100644
--- a/code/lib/store/src/csf/stepRunners.ts
+++ b/code/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();
   };
 }