Core: Fix decorator context update

examples/official-storybook/stories/core/decorators.stories.js

@@ -35,3 +35,23 @@ All.decorators = [
     </>
   ),
 ];
+
+export const PropOverride = (args, { prop1, prop2 }) => (
+  <p>
+    Story prop: <pre>{JSON.stringify({ prop1, prop2 })}</pre>
+  </p>
+);
+PropOverride.decorators = [
+  (Story) => (
+    <>
+      <p>First local Decorator</p>
+      <Story prop1="prop1" />
+    </>
+  ),
+  (Story) => (
+    <>
+      <p>Second Local Decorator</p>
+      <Story prop2="prop2" />
+    </>
+  ),
+];

lib/client-api/src/decorators.test.ts

@@ -42,6 +42,25 @@ describe('client-api.decorators', () => {
     expect(contexts.map((c) => c.k)).toEqual([0, 3, 2, 1]);
   });
 
+  it('passes context through to sub decorators additively', () => {
+    const contexts = [];
+    const decorators = [
+      (s, c) => contexts.push(c) && s({ b: 1 }),
+      (s, c) => contexts.push(c) && s({ c: 2 }),
+      (s, c) => contexts.push(c) && s({ d: 3 }),
+    ];
+    const decorated = defaultDecorateStory((c) => contexts.push(c), decorators);
+
+    expect(contexts).toEqual([]);
+    decorated(makeContext({ a: 0 }));
+    expect(contexts.map(({ a, b, c, d }) => ({ a, b, c, d }))).toEqual([
+      { a: 0, b: undefined, c: undefined, d: undefined },
+      { a: 0, b: undefined, c: undefined, d: 3 },
+      { a: 0, b: undefined, c: 2, d: 3 },
+      { a: 0, b: 1, c: 2, d: 3 },
+    ]);
+  });
+
   it('does not recreate decorated story functions each time', () => {
     const decoratedStories = [];
     const decorators = [

lib/client-api/src/decorators.ts

@@ -11,37 +11,21 @@ const defaultContext: StoryContext = {
   globals: {},
 };
 
-/**
- * When you call the story function inside a decorator, e.g.:
- *
- * ```jsx
- * <div>{storyFn({ foo: 'bar' })}</div>
- * ```
- *
- * This will override the `foo` property on the `innerContext`, which gets
- * merged in with the default context
- */
-const bindWithContext = (
-  storyFn: LegacyStoryFn,
-  getStoryContext: () => StoryContext
-): PartialStoryFn =>
-  // (NOTE: You cannot override the parameters key, it is fixed)
-  ({ id, name, kind, parameters, ...contextUpdate }: StoryContextUpdate = {}) =>
-    storyFn({ ...getStoryContext(), ...contextUpdate });
-
 export const decorateStory = (
   storyFn: LegacyStoryFn,
   decorator: DecoratorFunction,
-  getStoryContext: () => StoryContext
+  bindWithContext: (storyFn: LegacyStoryFn) => PartialStoryFn
 ): LegacyStoryFn => {
   // Bind the partially decorated storyFn so that when it is called it always knows about the story context,
   // no matter what it is passed directly. This is because we cannot guarantee a decorator will
   // pass the context down to the next decorated story in the chain.
-  const boundStoryFunction = bindWithContext(storyFn, getStoryContext);
+  const boundStoryFunction = bindWithContext(storyFn);
 
   return (context: StoryContext) => decorator(boundStoryFunction, context);
 };
 
+type ContextStore = { value: StoryContext };
+
 export const defaultDecorateStory = (
   storyFn: LegacyStoryFn,
   decorators: DecoratorFunction[]
@@ -52,13 +36,31 @@ export const defaultDecorateStory = (
   // (ie to this story), so there is no possibility of overlap.
   // This will break if you call the same story twice interleaved
   // (React might do it if you rendered the same story twice in the one ReactDom.render call, for instance)
-  let contextStore: StoryContext;
+  const contextStore: ContextStore = { value: defaultContext };
+
+  /**
+   * When you call the story function inside a decorator, e.g.:
+   *
+   * ```jsx
+   * <div>{storyFn({ foo: 'bar' })}</div>
+   * ```
+   *
+   * This will override the `foo` property on the `innerContext`, which gets
+   * merged in with the default context
+   */
+  const bindWithContext = (decoratedStoryFn: LegacyStoryFn): PartialStoryFn =>
+    // (NOTE: You cannot override the parameters key, it is fixed)
+    ({ id, name, kind, parameters, ...contextUpdate }: StoryContextUpdate = {}) => {
+      contextStore.value = { ...contextStore.value, ...contextUpdate };
+      return decoratedStoryFn(contextStore.value);
+    };
+
   const decoratedWithContextStore = decorators.reduce(
-    (story, decorator) => decorateStory(story, decorator, () => contextStore),
+    (story, decorator) => decorateStory(story, decorator, bindWithContext),
     storyFn
   );
   return (context = defaultContext) => {
-    contextStore = context;
+    contextStore.value = context;
     return decoratedWithContextStore(context); // Pass the context directly into the first decorator
   };
 };