From 646d8ccf17a2f322554eb6774d076303553a9360 Mon Sep 17 00:00:00 2001
From: Tom Coleman <tom@thesnail.org>
Date: Wed, 15 Sep 2021 11:45:09 +1000
Subject: [PATCH] Fix `useStories` to correctly respond to change in `storyId`.

---
 addons/docs/src/blocks/Source.tsx  |  6 +++---
 addons/docs/src/blocks/useStory.ts | 23 ++++++++++++-----------
 2 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/addons/docs/src/blocks/Source.tsx b/addons/docs/src/blocks/Source.tsx
index 44787d71eda..c720eb3b201 100644
--- a/addons/docs/src/blocks/Source.tsx
+++ b/addons/docs/src/blocks/Source.tsx
@@ -109,7 +109,7 @@ export const getSourceProps = (
   const targetIds = multiProps.ids || [targetId];
 
   const stories = useStories(targetIds, docsContext);
-  if (!stories) {
+  if (!stories.every(Boolean)) {
     return { error: SourceError.SOURCE_UNAVAILABLE, state: SourceState.NONE };
   }
 
@@ -117,13 +117,13 @@ export const getSourceProps = (
     source = targetIds
       .map((storyId, idx) => {
         const storySource = getStorySource(storyId, sourceContext);
-        const storyObj = stories[idx];
+        const storyObj = stories[idx] as Story;
         return getSnippet(storySource, storyObj);
       })
       .join('\n\n');
   }
 
-  const state = getSourceState(stories);
+  const state = getSourceState(stories as Story[]);
 
   const { docs: docsParameters = {} } = parameters;
   const { source: sourceParameters = {} } = docsParameters;
diff --git a/addons/docs/src/blocks/useStory.ts b/addons/docs/src/blocks/useStory.ts
index 6d11636eeb6..459abb13b27 100644
--- a/addons/docs/src/blocks/useStory.ts
+++ b/addons/docs/src/blocks/useStory.ts
@@ -4,7 +4,7 @@ import { Story } from '@storybook/store';
 
 import { DocsContextProps } from './DocsContext';
 
-export function useStory<TFramework extends AnyFramework>(
+export function useStory<TFramework extends AnyFramework = AnyFramework>(
   storyId: StoryId,
   context: DocsContextProps<TFramework>
 ): Story<TFramework> | void {
@@ -12,19 +12,20 @@ export function useStory<TFramework extends AnyFramework>(
   return stories && stories[0];
 }
 
-export function useStories<TFramework extends AnyFramework>(
+export function useStories<TFramework extends AnyFramework = AnyFramework>(
   storyIds: StoryId[],
   context: DocsContextProps<TFramework>
-): Story<TFramework>[] | void {
-  const [stories, setStories] = useState(null);
+): (Story<TFramework> | void)[] {
+  const [storiesById, setStories] = useState({} as Record<StoryId, Story<TFramework>>);
 
   useEffect(() => {
-    Promise.all(storyIds.map((storyId) => context.loadStory(storyId))).then((loadedStories) => {
-      if (!stories) {
-        setStories(loadedStories);
-      }
-    });
-  });
+    Promise.all(
+      storyIds.map(async (storyId) => {
+        const story = await context.loadStory(storyId);
+        setStories((current) => ({ ...current, [storyId]: story }));
+      })
+    );
+  }, storyIds);
 
-  return stories;
+  return storyIds.map((storyId) => storiesById[storyId]);
 }