Ensure we retained prepared-ness when the index HMRs

code/lib/api/src/lib/stories.ts

@@ -519,6 +519,21 @@ export const transformStoryIndexToStoriesHash = (
     .reduce(addItem, orphanHash);
 };
 
+export const addPreparedStories = (newHash: StoriesHash, oldHash?: StoriesHash) => {
+  if (!oldHash) return newHash;
+
+  return Object.fromEntries(
+    Object.entries(newHash).map(([id, newEntry]) => {
+      const oldEntry = oldHash[id];
+      if (newEntry.type === 'story' && oldEntry?.type === 'story' && oldEntry.prepared) {
+        return [id, { ...oldEntry, ...newEntry, prepared: true }];
+      }
+
+      return [id, newEntry];
+    })
+  );
+};
+
 export const getComponentLookupList = memoize(1)((hash: StoriesHash) => {
   return Object.entries(hash).reduce((acc, i) => {
     const value = i[1];

code/lib/api/src/modules/stories.ts

@@ -25,6 +25,7 @@ import {
   getStoriesLookupList,
   HashEntry,
   LeafEntry,
+  addPreparedStories,
 } from '../lib/stories';
 
 import type {
@@ -351,13 +352,16 @@ export const init: ModuleFn<SubAPI, SubState, true> = ({
       }
     },
     setStoryList: async (storyIndex: StoryIndex) => {
-      const hash = transformStoryIndexToStoriesHash(storyIndex, {
+      const newHash = transformStoryIndexToStoriesHash(storyIndex, {
         provider,
         docsOptions,
       });
 
+      // Now we need to patch in the existing prepared stories
+      const oldHash = store.getState().storiesHash;
+
       await store.setState({
-        storiesHash: hash,
+        storiesHash: addPreparedStories(newHash, oldHash),
         storiesConfigured: true,
         storiesFailed: null,
       });

code/lib/api/src/tests/stories.test.ts

@@ -1,4 +1,5 @@
-/// <reference types="jest" />;
+// Need to import jest as mockJest for annoying jest reasons. Is there a better way?
+import { jest, jest as mockJest, it, describe, expect, beforeEach } from '@jest/globals';
 
 import {
   STORY_ARGS_UPDATED,
@@ -21,17 +22,17 @@ import { StoryEntry, SetStoriesStoryData, SetStoriesStory, StoryIndex } from '..
 import type Store from '../store';
 import { ModuleArgs } from '..';
 
-const mockStories: jest.MockedFunction<() => StoryIndex['entries']> = jest.fn();
+const mockStories = jest.fn<StoryIndex['entries'], []>();
 
 jest.mock('../lib/events');
 jest.mock('global', () => ({
-  ...(jest.requireActual('global') as Record<string, any>),
-  fetch: jest.fn(() => ({ json: () => ({ v: 4, entries: mockStories() }) })),
+  ...(mockJest.requireActual('global') as Record<string, any>),
+  fetch: mockJest.fn(() => ({ json: () => ({ v: 4, entries: mockStories() }) })),
   FEATURES: { storyStoreV7: true },
   CONFIG_TYPE: 'DEVELOPMENT',
 }));
 
-const getEventMetadataMock = getEventMetadata as jest.MockedFunction<typeof getEventMetadata>;
+const getEventMetadataMock = getEventMetadata as ReturnType<typeof jest.fn>;
 
 const mockIndex = {
   'component-a--story-1': {
@@ -58,7 +59,7 @@ function createMockStore(initialState = {}) {
   let state = initialState;
   return {
     getState: jest.fn(() => state),
-    setState: jest.fn((s) => {
+    setState: jest.fn((s: typeof state) => {
       state = { ...state, ...s };
       return Promise.resolve(state);
     }),
@@ -1195,6 +1196,47 @@ describe('stories API', () => {
       expect(Object.keys(storedStoriesHash)).toEqual(['component-a', 'component-a--story-1']);
     });
 
+    it('retains prepared-ness of stories', async () => {
+      const navigate = jest.fn();
+      const store = createMockStore();
+      const fullAPI = Object.assign(new EventEmitter(), {
+        setStories: jest.fn(),
+        setOptions: jest.fn(),
+      });
+
+      const { api, init } = initStories({ store, navigate, provider, fullAPI } as any);
+      Object.assign(fullAPI, api);
+
+      global.fetch.mockClear();
+      await init();
+      expect(global.fetch).toHaveBeenCalledTimes(1);
+
+      fullAPI.emit(STORY_PREPARED, {
+        id: 'component-a--story-1',
+        parameters: { a: 'b' },
+        args: { c: 'd' },
+      });
+      // Let the promise/await chain resolve
+      await new Promise((r) => setTimeout(r, 0));
+      expect(store.getState().storiesHash['component-a--story-1'] as StoryEntry).toMatchObject({
+        prepared: true,
+        parameters: { a: 'b' },
+        args: { c: 'd' },
+      });
+
+      global.fetch.mockClear();
+      provider.serverChannel.emit(STORY_INDEX_INVALIDATED);
+      expect(global.fetch).toHaveBeenCalledTimes(1);
+
+      // Let the promise/await chain resolve
+      await new Promise((r) => setTimeout(r, 0));
+      expect(store.getState().storiesHash['component-a--story-1'] as StoryEntry).toMatchObject({
+        prepared: true,
+        parameters: { a: 'b' },
+        args: { c: 'd' },
+      });
+    });
+
     it('handles docs entries', async () => {
       mockStories.mockReset().mockReturnValue({
         'component-a--page': {