Merge pull request #22336 from redbugz/fix-hooks-rendered

React: Fix decorators to conditionally render children

code/lib/preview-api/src/modules/addons/hooks.ts

@@ -70,7 +70,7 @@ export class HooksContext<TRenderer extends Renderer, TArgs extends Args = Args>
   init() {
     this.hookListsMap = new WeakMap();
     this.mountedDecorators = new Set();
-    this.prevMountedDecorators = this.mountedDecorators;
+    this.prevMountedDecorators = new Set();
     this.currentHooks = [];
     this.nextHookIndex = 0;
     this.currentPhase = 'NONE';
@@ -191,7 +191,7 @@ export const applyHooks =
     );
     return (context) => {
       const { hooks } = context as { hooks: HooksContext<TRenderer> };
-      hooks.prevMountedDecorators = hooks.mountedDecorators;
+      hooks.prevMountedDecorators ??= new Set();
       hooks.mountedDecorators = new Set([storyFn, ...decorators]);
       hooks.currentContext = context;
       hooks.hasUpdates = false;

code/lib/preview-api/src/modules/store/hooks.test.ts

@@ -147,6 +147,25 @@ describe('Preview hooks', () => {
       run(story, [decorator]);
       expect(effect).toHaveBeenCalledTimes(1);
     });
+    it('handles decorator conditionally rendering the story', () => {
+      const effect = jest.fn();
+      const story = () => {
+        useEffect(effect, []);
+      };
+      const decorator = (storyFn: any) => {
+        const [counter, setCounter] = useState(0);
+        useEffect(() => {
+          setCounter((prevCounter) => prevCounter + 1);
+        }, [counter]);
+        if (counter % 2 === 1) storyFn();
+        return 'placeholder while waiting';
+      };
+      run(story, [decorator]);
+      run(story, [decorator]);
+      run(story, [decorator]);
+      run(story, [decorator]);
+      expect(effect).toHaveBeenCalledTimes(2);
+    });
     it('retriggers the effect if some of the deps are changed', () => {
       const effect = jest.fn();
       let counter = 0;

code/lib/preview-api/template/stories/decorators.stories.ts

@@ -1,7 +1,14 @@
 import { global as globalThis } from '@storybook/global';
-import type { PartialStoryFn, PlayFunctionContext, StoryContext } from '@storybook/types';
+import type {
+  ArgsStoryFn,
+  PartialStoryFn,
+  PlayFunctionContext,
+  StoryContext,
+} from '@storybook/types';
 import { within } from '@storybook/testing-library';
 import { expect } from '@storybook/jest';
+import { useEffect } from '@storybook/preview-api';
+import { STORY_ARGS_UPDATED, UPDATE_STORY_ARGS, RESET_STORY_ARGS } from '@storybook/core-events';
 
 export default {
   component: globalThis.Components.Pre,
@@ -25,3 +32,35 @@ export const Inheritance = {
     await expect(canvas.getByTestId('pre').innerText).toEqual('story component project starting');
   },
 };
+
+// NOTE this story is currently broken in Chromatic for both Vue2/Vue3
+// Issue: https://github.com/storybookjs/storybook/issues/22945
+export const Hooks = {
+  decorators: [
+    // decorator that uses hooks
+    (storyFn: PartialStoryFn, context: StoryContext) => {
+      useEffect(() => {});
+      return storyFn({ args: { ...context.args, text: `story ${context.args['text']}` } });
+    },
+    // conditional decorator, runs before the above
+    (storyFn: PartialStoryFn, context: StoryContext) =>
+      context.args.condition
+        ? storyFn()
+        : (context.originalStoryFn as ArgsStoryFn)(context.args, context),
+  ],
+  args: {
+    text: 'text',
+    condition: true,
+  },
+  play: async ({ id, args }: PlayFunctionContext<any>) => {
+    const channel = globalThis.__STORYBOOK_ADDONS_CHANNEL__;
+    await channel.emit(RESET_STORY_ARGS, { storyId: id });
+    await new Promise((resolve) => channel.once(STORY_ARGS_UPDATED, resolve));
+
+    await channel.emit(UPDATE_STORY_ARGS, {
+      storyId: id,
+      updatedArgs: { condition: false },
+    });
+    await new Promise((resolve) => channel.once(STORY_ARGS_UPDATED, resolve));
+  },
+};