From 4d1d54af26175a87edabc80ffb69bdbaabcbfd1b Mon Sep 17 00:00:00 2001
From: Kasper Peulen <kasperpeulen@gmail.com>
Date: Thu, 7 Nov 2024 16:55:27 +0100
Subject: [PATCH 1/2] Add the afterEach hook

---
 src/story.test.ts |  8 ++++++++
 src/story.ts      | 15 +++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/src/story.test.ts b/src/story.test.ts
index f49c18f..8948e3c 100644
--- a/src/story.test.ts
+++ b/src/story.test.ts
@@ -39,6 +39,8 @@ async function doSomething() {
   a = 2;
 }
 
+async function validateSomething() {}
+
 async function cleanup() {
   a = 1;
 }
@@ -55,6 +57,9 @@ const simple: XMeta = {
     await doSomething();
     return cleanup;
   },
+  async afterEach() {
+    await validateSomething();
+  },
   args: { x: '1' },
   argTypes: { x: { type: { name: 'string' } } },
 };
@@ -71,6 +76,9 @@ const strict: XMeta<ButtonArgs> = {
     await doSomething();
     return cleanup;
   },
+  async afterEach() {
+    await validateSomething();
+  },
   argTypes: { x: { type: { name: 'string' } } },
 };
 
diff --git a/src/story.ts b/src/story.ts
index 4bcf3c7..86eae2e 100644
--- a/src/story.ts
+++ b/src/story.ts
@@ -263,6 +263,10 @@ export type BeforeEach<TRenderer extends Renderer = Renderer, TArgs = Args> = (
   context: StoryContext<TRenderer, TArgs>
 ) => Awaitable<CleanupCallback | void>;
 
+export type AfterEach<TRenderer extends Renderer = Renderer, TArgs = Args> = (
+  context: StoryContext<TRenderer, TArgs>
+) => Awaitable<void>;
+
 export interface Canvas {}
 
 export interface StoryContext<TRenderer extends Renderer = Renderer, TArgs = Args>
@@ -381,6 +385,17 @@ export interface BaseAnnotations<TRenderer extends Renderer = Renderer, TArgs =
    */
   beforeEach?: BeforeEach<TRenderer, TArgs>[] | BeforeEach<TRenderer, TArgs>;
 
+  /**
+   * Function to be called after each play function for post-test assertions.
+   * Don't use this function for cleaning up state.
+   * You can use the return callback of `beforeEach` for that, which is run when switching stories.
+   * When the function is async, it will be awaited.
+   *
+   * `afterEach` can be added to preview, the default export and to a specific story.
+   * They are run (and awaited) reverse order: preview, default export, story
+   */
+  afterEach?: AfterEach<TRenderer, TArgs>[] | BeforeEach<TRenderer, TArgs>;
+
   /**
    * Define a custom render function for the story(ies). If not passed, a default render function by the renderer will be used.
    */

From 1cc995764cc2f11eed3943250059a60bf0d84670 Mon Sep 17 00:00:00 2001
From: Kasper Peulen <kasperpeulen@gmail.com>
Date: Tue, 12 Nov 2024 15:30:26 +0100
Subject: [PATCH 2/2] Fix typo

---
 src/story.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/story.ts b/src/story.ts
index 86eae2e..7560bfe 100644
--- a/src/story.ts
+++ b/src/story.ts
@@ -394,7 +394,7 @@ export interface BaseAnnotations<TRenderer extends Renderer = Renderer, TArgs =
    * `afterEach` can be added to preview, the default export and to a specific story.
    * They are run (and awaited) reverse order: preview, default export, story
    */
-  afterEach?: AfterEach<TRenderer, TArgs>[] | BeforeEach<TRenderer, TArgs>;
+  afterEach?: AfterEach<TRenderer, TArgs>[] | AfterEach<TRenderer, TArgs>;
 
   /**
    * Define a custom render function for the story(ies). If not passed, a default render function by the renderer will be used.