Preview hooks: trigger effects after story render

app/angular/src/client/preview/angular/decorators.test.ts

@@ -1,3 +1,5 @@
+import addons, { mockChannel } from '@storybook/addons';
+
 import { moduleMetadata } from './decorators';
 import { addDecorator, storiesOf, clearDecorators, getStorybook } from '..';
 
@@ -78,6 +80,7 @@ describe('moduleMetadata', () => {
       imports: [MockModule],
     };
 
+    addons.setChannel(mockChannel());
     addDecorator(moduleMetadata(metadata));
 
     storiesOf('Test', module).add('Default', () => ({

app/react/src/client/preview/render.tsx

@@ -7,14 +7,16 @@ import { RenderMainArgs } from './types';
 
 const rootEl = document ? document.getElementById('root') : null;
 
-function render(node: React.ReactElement, el: Element) {
-  ReactDOM.render(
-    process.env.STORYBOOK_EXAMPLE_APP ? <React.StrictMode>{node}</React.StrictMode> : node,
-    el
-  );
-}
+const render = (node: React.ReactElement, el: Element) =>
+  new Promise(resolve => {
+    ReactDOM.render(
+      process.env.STORYBOOK_EXAMPLE_APP ? <React.StrictMode>{node}</React.StrictMode> : node,
+      el,
+      resolve
+    );
+  });
 
-export default function renderMain({
+export default async function renderMain({
   storyFn: StoryFn,
   selectedKind,
   selectedStory,
@@ -55,6 +57,6 @@ export default function renderMain({
     ReactDOM.unmountComponentAtNode(rootEl);
   }
 
-  render(element, rootEl);
+  await render(element, rootEl);
   showMain();
 }

examples/html-kitchen-sink/stories/button.stories.js

@@ -1,5 +1,6 @@
 import { document } from 'global';
 import { action } from '@storybook/addon-actions';
+import { useEffect } from '@storybook/client-api';
 
 export default {
   title: 'Demo',
@@ -15,3 +16,11 @@ export const button = () => {
   btn.addEventListener('click', action('Click'));
   return btn;
 };
+
+export const effect = () => {
+  useEffect(() => {
+    document.getElementById('button').style.backgroundColor = 'yellow';
+  });
+
+  return '<button id="button">I should be yellow</button>';
+};

examples/html-kitchen-sink/tests/__snapshots__/htmlshots.test.js.snap

@@ -215,6 +215,14 @@ exports[`Storyshots Demo button 1`] = `
 </button>
 `;
 
+exports[`Storyshots Demo effect 1`] = `
+<button
+  id="button"
+>
+  I should be yellow
+</button>
+`;
+
 exports[`Storyshots Demo heading 1`] = `
 <h1>
   Hello World

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

@@ -1,12 +1,13 @@
 /* eslint-disable no-underscore-dangle */
 import { logger } from '@storybook/client-logger';
-import { mockChannel } from '@storybook/addons';
+import addons, { mockChannel } from '@storybook/addons';
 import ClientApi from './client_api';
 import ConfigApi from './config_api';
 import StoryStore from './story_store';
 
 export const getContext = (() => decorateStory => {
   const channel = mockChannel();
+  addons.setChannel(channel);
   const storyStore = new StoryStore({ channel });
   const clientApi = new ClientApi({ storyStore, decorateStory });
   const { clearDecorators } = clientApi;

lib/client-api/src/hooks.test.js

@@ -1,4 +1,4 @@
-import { FORCE_RE_RENDER } from '@storybook/addon-contexts/src/shared/constants';
+import { FORCE_RE_RENDER, STORY_RENDERED } from '@storybook/core-events';
 import { defaultDecorateStory } from './client_api';
 import {
   applyHooks,
@@ -22,6 +22,11 @@ beforeEach(() => {
   mockChannel = {
     emit: jest.fn(),
     on: jest.fn(),
+    once: (event, callback) => {
+      if (event === STORY_RENDERED) {
+        callback();
+      }
+    },
     removeListener: jest.fn(),
   };
 });
@@ -359,6 +364,19 @@ describe('Preview hooks', () => {
       expect(storyFn).toHaveBeenCalledTimes(2);
       expect(state).toBe('bar');
     });
+    it('triggers only the last effect when updating state synchronously', () => {
+      const effects = [jest.fn(), jest.fn()];
+      const storyFn = jest.fn(() => {
+        const [effectIndex, setEffectIndex] = useState(0);
+        useEffect(effects[effectIndex], [effectIndex]);
+        if (effectIndex === 0) {
+          setEffectIndex(1);
+        }
+      });
+      run(storyFn);
+      expect(effects[0]).not.toHaveBeenCalled();
+      expect(effects[1]).toHaveBeenCalledTimes(1);
+    });
     it('performs synchronous state updates with updater function', () => {
       let state;
       let setState;

lib/client-api/src/hooks.ts

@@ -1,6 +1,6 @@
 import { logger } from '@storybook/client-logger';
 import addons, { StoryGetter, StoryContext } from '@storybook/addons';
-import { FORCE_RE_RENDER } from '@storybook/core-events';
+import { FORCE_RE_RENDER, STORY_RENDERED } from '@storybook/core-events';
 
 interface StoryStore {
   fromId: (
@@ -32,6 +32,7 @@ type AbstractFunction = (...args: any[]) => any;
 
 const hookListsMap = new WeakMap<AbstractFunction, Hook[]>();
 let mountedDecorators = new Set<AbstractFunction>();
+let prevMountedDecorators = mountedDecorators;
 
 let currentHooks: Hook[] = [];
 let nextHookIndex = 0;
@@ -72,7 +73,7 @@ const hookify = (fn: AbstractFunction) => (...args: any[]) => {
   const prevDecoratorName = currentDecoratorName;
 
   currentDecoratorName = fn.name;
-  if (mountedDecorators.has(fn)) {
+  if (prevMountedDecorators.has(fn)) {
     currentPhase = 'UPDATE';
     currentHooks = hookListsMap.get(fn) || [];
   } else {
@@ -105,13 +106,16 @@ export const applyHooks = (
 ) => (getStory: StoryGetter, decorators: Decorator[]) => {
   const decorated = applyDecorators(hookify(getStory), decorators.map(hookify));
   return (context: StoryContext) => {
+    prevMountedDecorators = mountedDecorators;
+    mountedDecorators = new Set([getStory, ...decorators]);
     currentContext = context;
     hasUpdates = false;
     let result = decorated(context);
-    mountedDecorators = new Set([getStory, ...decorators]);
     numberOfRenders = 1;
     while (hasUpdates) {
       hasUpdates = false;
+      currentEffects = [];
+      prevMountedDecorators = mountedDecorators;
       result = decorated(context);
       numberOfRenders += 1;
       if (numberOfRenders > RENDER_LIMIT) {
@@ -120,8 +124,10 @@ export const applyHooks = (
         );
       }
     }
-    triggerEffects();
-    currentContext = null;
+    addons.getChannel().once(STORY_RENDERED, () => {
+      triggerEffects();
+      currentContext = null;
+    });
     return result;
   };
 };
@@ -271,7 +277,7 @@ export function useReducer<S, A>(
 
 /*
   Triggers a side effect, see https://reactjs.org/docs/hooks-reference.html#usestate
-  Effects are triggered synchronously after calling the decorated story function
+  Effects are triggered synchronously after rendering the story
 */
 export function useEffect(create: () => (() => void) | void, deps?: any[]): void {
   const effect = useMemoLike('useEffect', () => ({ create }), deps);

lib/core/src/client/preview/start.js

@@ -226,8 +226,10 @@ export default function start(render, { decorateStory } = {}) {
       case 'story':
       default: {
         if (getDecorated) {
-          render(renderContext);
-          addons.getChannel().emit(Events.STORY_RENDERED, id);
+          (async () => {
+            await render(renderContext);
+            addons.getChannel().emit(Events.STORY_RENDERED, id);
+          })();
         } else {
           showNopreview();
           addons.getChannel().emit(Events.STORY_MISSING, id);