Allow using hooks after calling story function

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

@@ -37,6 +37,17 @@ describe('Preview hooks', () => {
       ]);
       expect(effect).toHaveBeenCalledTimes(1);
     });
+    it('triggers the effect from decorator if story call comes before useEffect', () => {
+      const effect = jest.fn();
+      run(() => {}, [
+        storyFn => {
+          const story = storyFn();
+          useEffect(effect);
+          return story;
+        },
+      ]);
+      expect(effect).toHaveBeenCalledTimes(1);
+    });
     it('retriggers the effect if no deps array is provided', () => {
       const effect = jest.fn();
       const storyFn = () => {
@@ -55,6 +66,27 @@ describe('Preview hooks', () => {
       run(storyFn);
       expect(effect).toHaveBeenCalledTimes(1);
     });
+    it("doesn't retrigger the effect from decorator if story has changed", () => {
+      const effect = jest.fn();
+      const decorator = storyFn => {
+        useEffect(effect, []);
+        return storyFn();
+      };
+      run(() => {}, [decorator]);
+      run(() => {}, [decorator]);
+      expect(effect).toHaveBeenCalledTimes(1);
+    });
+    it("doesn't retrigger the effect from decorator if story has changed and story call comes before useEffect", () => {
+      const effect = jest.fn();
+      const decorator = storyFn => {
+        const story = storyFn();
+        useEffect(effect, []);
+        return story;
+      };
+      run(() => {}, [decorator]);
+      run(() => {}, [decorator]);
+      expect(effect).toHaveBeenCalledTimes(1);
+    });
     it('retriggers the effect if some of the deps are changed', () => {
       const effect = jest.fn();
       let counter = 0;

lib/client-api/src/hooks.ts

@@ -51,6 +51,11 @@ const triggerEffects = () => {
 };
 
 const hookify = (fn: AbstractFunction) => (...args: any[]) => {
+  const prevPhase = currentPhase;
+  const prevHooks = currentHooks;
+  const prevNextHookIndex = nextHookIndex;
+  const prevDecoratorName = currentDecoratorName;
+
   currentDecoratorName = fn.name;
   if (mountedDecorators.has(fn)) {
     currentPhase = 'UPDATE';
@@ -70,7 +75,10 @@ const hookify = (fn: AbstractFunction) => (...args: any[]) => {
     );
   }
 
-  currentDecoratorName = null;
+  currentPhase = prevPhase;
+  currentHooks = prevHooks;
+  nextHookIndex = prevNextHookIndex;
+  currentDecoratorName = prevDecoratorName;
   return result;
 };
 
@@ -96,7 +104,6 @@ export const applyHooks = (
         );
       }
     }
-    currentPhase = 'NONE';
     triggerEffects();
     return result;
   };