Fix useEfect in inline Docs

examples/official-storybook/stories/hooks.stories.js

@@ -1,5 +1,5 @@
 import React from 'react';
-import { useState } from '@storybook/client-api';
+import { useEffect, useRef, useState } from '@storybook/client-api';
 
 export default {
   title: 'Hooks',
@@ -19,3 +19,18 @@ export const Input = () => {
   const [text, setText] = useState('foo');
   return <input value={text} onChange={e => setText(e.target.value)} />;
 };
+
+export const effect = () => {
+  const ref = useRef();
+  useEffect(() => {
+    if (ref.current != null) {
+      ref.current.style.backgroundColor = 'yellow';
+    }
+  });
+
+  return (
+    <button type="button" ref={ref}>
+      I should be yellow
+    </button>
+  );
+};

lib/addons/src/hooks.ts

@@ -1,6 +1,6 @@
 import { window } from 'global';
 import { logger } from '@storybook/client-logger';
-import { FORCE_RE_RENDER, STORY_RENDERED } from '@storybook/core-events';
+import { FORCE_RE_RENDER, STORY_RENDERED, DOCS_RENDERED } from '@storybook/core-events';
 import addons, { StoryGetter, StoryContext } from './public_api';
 
 interface StoryStore {
@@ -31,6 +31,8 @@ interface Effect {
 type Decorator = (getStory: StoryGetter, context: StoryContext) => any;
 type AbstractFunction = (...args: any[]) => any;
 
+const RenderEvents = [STORY_RENDERED, DOCS_RENDERED];
+
 export class HooksContext {
   hookListsMap: WeakMap<AbstractFunction, Hook[]>;
 
@@ -54,6 +56,12 @@ export class HooksContext {
 
   currentContext: StoryContext | null;
 
+  renderListener = () => {
+    this.triggerEffects();
+    this.currentContext = null;
+    this.removeRenderListeners();
+  };
+
   constructor() {
     this.init();
   }
@@ -79,6 +87,7 @@ export class HooksContext {
       }
     });
     this.init();
+    this.removeRenderListeners();
   }
 
   getNextHook() {
@@ -104,6 +113,16 @@ export class HooksContext {
     this.prevEffects = this.currentEffects;
     this.currentEffects = [];
   }
+
+  addRenderListeners() {
+    const channel = addons.getChannel();
+    RenderEvents.forEach(e => channel.on(e, this.renderListener));
+  }
+
+  removeRenderListeners() {
+    const channel = addons.getChannel();
+    RenderEvents.forEach(e => channel.removeListener(e, this.renderListener));
+  }
 }
 
 const hookify = (fn: AbstractFunction) => (...args: any[]) => {
@@ -170,10 +189,7 @@ export const applyHooks = (
         );
       }
     }
-    addons.getChannel().once(STORY_RENDERED, () => {
-      hooks.triggerEffects();
-      hooks.currentContext = null;
-    });
+    hooks.addRenderListeners();
     return result;
   };
 };

lib/core-events/src/index.ts

@@ -18,6 +18,7 @@ enum events {
   STORY_MISSING = 'storyMissing',
   STORY_ERRORED = 'storyErrored',
   STORY_THREW_EXCEPTION = 'storyThrewException',
+  DOCS_RENDERED = 'docsRendered',
 }
 
 // Enables: `import Events from ...`
@@ -25,21 +26,24 @@ export default events;
 
 // Enables: `import * as Events from ...` or `import { CHANNEL_CREATED } as Events from ...`
 // This is the preferred method
-export const { CHANNEL_CREATED } = events;
-export const { GET_CURRENT_STORY } = events;
-export const { SET_CURRENT_STORY } = events;
-export const { GET_STORIES } = events;
-export const { SET_STORIES } = events;
-export const { STORIES_CONFIGURED } = events;
-export const { SELECT_STORY } = events;
-export const { PREVIEW_KEYDOWN } = events;
-export const { FORCE_RE_RENDER } = events;
-export const { REGISTER_SUBSCRIPTION } = events;
-export const { STORY_INIT } = events;
-export const { STORY_ADDED } = events;
-export const { STORY_RENDER } = events;
-export const { STORY_RENDERED } = events;
-export const { STORY_MISSING } = events;
-export const { STORY_ERRORED } = events;
-export const { STORY_CHANGED } = events;
-export const { STORY_THREW_EXCEPTION } = events;
+export const {
+  CHANNEL_CREATED,
+  GET_CURRENT_STORY,
+  SET_CURRENT_STORY,
+  GET_STORIES,
+  SET_STORIES,
+  STORIES_CONFIGURED,
+  SELECT_STORY,
+  PREVIEW_KEYDOWN,
+  FORCE_RE_RENDER,
+  REGISTER_SUBSCRIPTION,
+  STORY_INIT,
+  STORY_ADDED,
+  STORY_RENDER,
+  STORY_RENDERED,
+  STORY_MISSING,
+  STORY_ERRORED,
+  STORY_CHANGED,
+  STORY_THREW_EXCEPTION,
+  DOCS_RENDERED,
+} = events;

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

@@ -239,7 +239,8 @@ export default function start(render, { decorateStory } = {}) {
           <DocsContainer context={renderContext}>
             <Page />
           </DocsContainer>,
-          document.getElementById('docs-root')
+          document.getElementById('docs-root'),
+          () => addons.getChannel().emit(Events.DOCS_RENDERED, kind)
         );
         break;
       }