Merge pull request #7992 from storybookjs/fix-docs-effects

API: Fix useEfect in inline Docs
storybookjs · Sep 4, 2019 · b41b3ab · b41b3ab
2 parents bdf6f22 + 6be0539
commit b41b3ab
Show file tree

Hide file tree

Showing 5 changed files with 87 additions and 37 deletions.
diff --git a/examples/official-storybook/stories/hooks.stories.js b/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',
@@ -20,6 +20,21 @@ export const Input = () => {
   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>
+  );
+};
+
 export const reactHookCheckbox = () => {
   const [on, setOn] = React.useState(false);
   return (

diff --git a/lib/addons/src/hooks.ts b/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;
   };
 };

diff --git a/lib/client-api/src/hooks.test.js b/lib/client-api/src/hooks.test.js
@@ -19,18 +19,32 @@ jest.mock('@storybook/client-logger', () => ({
   logger: { warn: jest.fn(), log: jest.fn() },
 }));
 
+const SOME_EVENT = 'someEvent';
 let mockChannel;
 let hooks;
+let onSomeEvent;
+let removeSomeEventListener;
 beforeEach(() => {
+  onSomeEvent = jest.fn();
+  removeSomeEventListener = jest.fn();
   mockChannel = {
     emit: jest.fn(),
-    on: jest.fn(),
-    once: (event, callback) => {
-      if (event === STORY_RENDERED) {
-        callback();
+    on(event, callback) {
+      switch (event) {
+        case STORY_RENDERED:
+          callback();
+          break;
+        case SOME_EVENT:
+          onSomeEvent(event, callback);
+          break;
+        default:
+      }
+    },
+    removeListener(event, callback) {
+      if (event === SOME_EVENT) {
+        removeSomeEventListener(event, callback);
       }
     },
-    removeListener: jest.fn(),
   };
   hooks = new HooksContext();
   addons.setChannel(mockChannel);
@@ -166,28 +180,28 @@ describe('Preview hooks', () => {
       run(() => {}, [
         storyFn => {
           useChannel({
-            SOME_EVENT: handler,
+            [SOME_EVENT]: handler,
           });
           return storyFn();
         },
       ]);
-      expect(mockChannel.on).toHaveBeenCalledTimes(1);
-      expect(mockChannel.removeListener).toHaveBeenCalledTimes(0);
+      expect(onSomeEvent).toHaveBeenCalledTimes(1);
+      expect(removeSomeEventListener).toHaveBeenCalledTimes(0);
     });
     it('calls .removeListener when removing the decorator', () => {
       const handler = () => {};
       run(() => {}, [
         storyFn => {
           useChannel({
-            SOME_EVENT: handler,
+            [SOME_EVENT]: handler,
           });
           return storyFn();
         },
       ]);
-      expect(mockChannel.on).toHaveBeenCalledTimes(1);
-      expect(mockChannel.removeListener).toHaveBeenCalledTimes(0);
+      expect(onSomeEvent).toHaveBeenCalledTimes(1);
+      expect(removeSomeEventListener).toHaveBeenCalledTimes(0);
       run(() => {});
-      expect(mockChannel.removeListener).toHaveBeenCalledTimes(1);
+      expect(removeSomeEventListener).toHaveBeenCalledTimes(1);
     });
   });
   describe('useStoryContext', () => {

diff --git a/lib/core-events/src/index.ts b/lib/core-events/src/index.ts
@@ -18,28 +18,32 @@ enum events {
   STORY_MISSING = 'storyMissing',
   STORY_ERRORED = 'storyErrored',
   STORY_THREW_EXCEPTION = 'storyThrewException',
+  DOCS_RENDERED = 'docsRendered',
 }
 
 // Enables: `import Events from ...`
 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;
diff --git a/lib/core/src/client/preview/start.js b/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;
       }