From fc9853ace3057ac9e98432b20b3610d1c0e22dd0 Mon Sep 17 00:00:00 2001
From: Michael Shilman <shilman@users.noreply.github.com>
Date: Thu, 9 Dec 2021 09:19:40 +0800
Subject: [PATCH] Merge pull request #16947 from
 storybookjs/16877-await-init-before-cache

Core: Ensure we have a full story index before caching
---
 lib/store/src/StoryStore.test.ts | 18 +++++++++++++++++-
 lib/store/src/StoryStore.ts      | 10 ++++++----
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/lib/store/src/StoryStore.test.ts b/lib/store/src/StoryStore.test.ts
index 1a85a832c686..8b882b373f9b 100644
--- a/lib/store/src/StoryStore.test.ts
+++ b/lib/store/src/StoryStore.test.ts
@@ -281,7 +281,7 @@ describe('StoryStore', () => {
       const store = new StoryStore();
       store.setProjectAnnotations(projectAnnotations);
       store.initialize({ storyIndex, importFn, cache: false });
-      await store.cacheAllCSFFiles(false);
+      await store.cacheAllCSFFiles();
 
       await store.loadStory({ storyId: 'component-one--a' });
       expect(importFn).toHaveBeenCalledWith(storyIndex.stories['component-one--a'].importPath);
@@ -943,4 +943,20 @@ describe('StoryStore', () => {
       });
     });
   });
+
+  describe('cacheAllCsfFiles', () => {
+    describe('if the store is not yet initialized', () => {
+      it('waits for initialization', async () => {
+        const store = new StoryStore();
+
+        importFn.mockClear();
+        const cachePromise = store.cacheAllCSFFiles();
+
+        store.setProjectAnnotations(projectAnnotations);
+        store.initialize({ storyIndex, importFn, cache: false });
+
+        await expect(cachePromise).resolves.toEqual(undefined);
+      });
+    });
+  });
 });
diff --git a/lib/store/src/StoryStore.ts b/lib/store/src/StoryStore.ts
index b4d486b47c70..ddfeb57c14c4 100644
--- a/lib/store/src/StoryStore.ts
+++ b/lib/store/src/StoryStore.ts
@@ -81,7 +81,7 @@ export class StoryStore<TFramework extends AnyFramework> {
 
   prepareStoryWithCache: typeof prepareStory;
 
-  initializationPromise: Promise<void>;
+  initializationPromise: SynchronousPromise<void>;
 
   resolveInitializationPromise: () => void;
 
@@ -175,9 +175,11 @@ export class StoryStore<TFramework extends AnyFramework> {
   }
 
   cacheAllCSFFiles(): PromiseLike<void> {
-    return this.loadAllCSFFiles().then((csfFiles) => {
-      this.cachedCSFFiles = csfFiles;
-    });
+    return this.initializationPromise.then(() =>
+      this.loadAllCSFFiles().then((csfFiles) => {
+        this.cachedCSFFiles = csfFiles;
+      })
+    );
   }
 
   // Load the CSF file for a story and prepare the story from it and the project annotations.