From 3e86b57aba5586d5b20f55ee555cf6ed9b65d84d Mon Sep 17 00:00:00 2001
From: Gert Hengeveld <info@ghengeveld.nl>
Date: Mon, 17 Jun 2024 13:53:44 +0200
Subject: [PATCH] Add support for beforeAll hook to projectAnnotations

---
 .../src/modules/preview-web/Preview.tsx       | 32 ++++++++++++++-----
 .../src/modules/store/csf/composeConfigs.ts   |  1 +
 code/lib/types/src/modules/story.ts           |  2 ++
 3 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/code/lib/preview-api/src/modules/preview-web/Preview.tsx b/code/lib/preview-api/src/modules/preview-web/Preview.tsx
index ee4524c32ecf..f578fcc1893d 100644
--- a/code/lib/preview-api/src/modules/preview-web/Preview.tsx
+++ b/code/lib/preview-api/src/modules/preview-web/Preview.tsx
@@ -72,6 +72,8 @@ export class Preview<TRenderer extends Renderer> {
   // project annotations. Once the index loads, it is stored on the store and this will get unset.
   private projectAnnotationsBeforeInitialization?: ProjectAnnotations<TRenderer>;
 
+  private beforeAllCleanup?: () => MaybePromise<void>;
+
   protected storeInitializationPromise: Promise<void>;
 
   protected resolveStoreInitializationPromise!: () => void;
@@ -160,9 +162,24 @@ export class Preview<TRenderer extends Renderer> {
   }
 
   // If initialization gets as far as project annotations, this function runs.
-  async initializeWithProjectAnnotations(projectAnnotations: ProjectAnnotations<TRenderer>) {
-    this.projectAnnotationsBeforeInitialization = projectAnnotations;
+  async initializeWithProjectAnnotations(
+    projectAnnotations: ProjectAnnotations<TRenderer>,
+    storyStore?: StoryStore<TRenderer>
+  ) {
     try {
+      await this.beforeAllCleanup?.();
+      const cleanup = await projectAnnotations.beforeAll?.();
+      if (cleanup) this.beforeAllCleanup = cleanup;
+    } catch (err) {
+      this.renderPreviewEntryError('Error in beforeAll hook:', err as Error);
+      throw err;
+    }
+
+    // Don't reinitialize story store if it's already been set.
+    if (storyStore) return;
+
+    try {
+      this.projectAnnotationsBeforeInitialization = projectAnnotations;
       const storyIndex = await this.getStoryIndexFromServer();
       return this.initializeWithStoryIndex(storyIndex);
     } catch (err) {
@@ -226,13 +243,12 @@ export class Preview<TRenderer extends Renderer> {
     this.getProjectAnnotations = getProjectAnnotations;
 
     const projectAnnotations = await this.getProjectAnnotationsOrRenderError();
-    if (!this.storyStoreValue) {
-      await this.initializeWithProjectAnnotations(projectAnnotations);
-      return;
-    }
+    await this.initializeWithProjectAnnotations(projectAnnotations, this.storyStoreValue);
 
-    this.storyStoreValue.setProjectAnnotations(projectAnnotations);
-    this.emitGlobals();
+    if (this.storyStoreValue) {
+      this.storyStoreValue.setProjectAnnotations(projectAnnotations);
+      this.emitGlobals();
+    }
   }
 
   async onStoryIndexChanged() {
diff --git a/code/lib/preview-api/src/modules/store/csf/composeConfigs.ts b/code/lib/preview-api/src/modules/store/csf/composeConfigs.ts
index b6ed3ecd7551..5eb1734093e8 100644
--- a/code/lib/preview-api/src/modules/store/csf/composeConfigs.ts
+++ b/code/lib/preview-api/src/modules/store/csf/composeConfigs.ts
@@ -59,6 +59,7 @@ export function composeConfigs<TRenderer extends Renderer>(
     initialGlobals: getObjectField(moduleExportList, 'initialGlobals'),
     globalTypes: getObjectField(moduleExportList, 'globalTypes'),
     loaders: getArrayField(moduleExportList, 'loaders'),
+    beforeAll: getSingletonField(moduleExportList, 'beforeAll'),
     beforeEach: getArrayField(moduleExportList, 'beforeEach'),
     render: getSingletonField(moduleExportList, 'render'),
     renderToCanvas: getSingletonField(moduleExportList, 'renderToCanvas'),
diff --git a/code/lib/types/src/modules/story.ts b/code/lib/types/src/modules/story.ts
index 041aaed2e709..4bbfce1cf4f4 100644
--- a/code/lib/types/src/modules/story.ts
+++ b/code/lib/types/src/modules/story.ts
@@ -46,6 +46,8 @@ export type ProjectAnnotations<TRenderer extends Renderer> = CsfProjectAnnotatio
 
   /* @deprecated use renderToCanvas */
   renderToDOM?: RenderToCanvas<TRenderer>;
+
+  beforeAll?: () => MaybePromise<void | (() => MaybePromise<void>)>;
 };
 
 type NamedExportsOrDefault<TExport> = TExport | { default: TExport };