From e371cb4bf8b899ad30baf4195204e408d2f6003f Mon Sep 17 00:00:00 2001
From: Phil Pluckthun <phil@kitten.sh>
Date: Tue, 4 Jun 2019 17:35:39 +0100
Subject: [PATCH] Simplify useImmediateEffect and add additional test

---
 src/hooks/useImmediateEffect.test.ts | 33 +++++++++++++++++++++++-----
 src/hooks/useImmediateEffect.ts      | 11 +++++-----
 2 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/src/hooks/useImmediateEffect.test.ts b/src/hooks/useImmediateEffect.test.ts
index 8e9c3fb77f..0f2df4c932 100644
--- a/src/hooks/useImmediateEffect.test.ts
+++ b/src/hooks/useImmediateEffect.test.ts
@@ -1,5 +1,5 @@
 import React from 'react';
-import { renderHook } from 'react-hooks-testing-library';
+import { renderHook, act } from 'react-hooks-testing-library';
 import { useImmediateEffect } from './useImmediateEffect';
 
 it('calls effects immediately on mount', () => {
@@ -8,14 +8,37 @@ it('calls effects immediately on mount', () => {
   const effect = jest.fn();
 
   spy.mockImplementation(useEffect);
-  renderHook(() => useImmediateEffect(effect, [effect]));
 
-  expect(effect).toHaveBeenCalledTimes(1);
-  expect(effect).toHaveBeenCalledTimes(1);
+  renderHook(() => {
+    useImmediateEffect(effect, [effect]);
+    expect(effect).toHaveBeenCalledTimes(1);
+  });
 
+  expect(effect).toHaveBeenCalledTimes(1);
   expect(useEffect).toHaveBeenCalledWith(expect.any(Function), [
     expect.any(Function),
   ]);
 
-  useEffect.mockRestore();
+  spy.mockRestore();
+});
+
+it('behaves like useEffect otherwise', () => {
+  const spy = jest.spyOn(React, 'useEffect');
+  const effect = jest.fn();
+
+  const { result } = renderHook(() => {
+    const [ref, setState] = React.useState({});
+    useImmediateEffect(effect, [ref]);
+
+    expect(effect).toHaveBeenCalledTimes(1);
+    expect(spy).toHaveBeenCalled();
+
+    const forceUpdate = () => setState({});
+    return forceUpdate;
+  });
+
+  act(() => result.current()); // forceUpdate
+  expect(effect).toHaveBeenCalledTimes(2);
+
+  spy.mockRestore();
 });
diff --git a/src/hooks/useImmediateEffect.ts b/src/hooks/useImmediateEffect.ts
index 8f6e25edb7..df76416788 100644
--- a/src/hooks/useImmediateEffect.ts
+++ b/src/hooks/useImmediateEffect.ts
@@ -1,4 +1,6 @@
-import { useRef, useEffect, useCallback } from 'react';
+/* eslint-disable react-hooks/exhaustive-deps */
+
+import { useRef, useEffect } from 'react';
 import { noop } from '../utils';
 
 enum LifecycleState {
@@ -16,22 +18,21 @@ export const useImmediateEffect = (
 ) => {
   const teardown = useRef(noop);
   const state = useRef(LifecycleState.WillMount);
-  const execute = useCallback(effect, changes);
 
   // On initial render we just execute the effect
   if (state.current === LifecycleState.WillMount) {
     state.current = LifecycleState.DidMount;
-    teardown.current = execute();
+    teardown.current = effect();
   }
 
   useEffect(() => {
     // Initially we skip executing the effect since we've already done so on
     // initial render, then we execute it as usual
     if (state.current === LifecycleState.Update) {
-      return (teardown.current = execute());
+      return (teardown.current = effect());
     } else {
       state.current = LifecycleState.Update;
       return teardown.current;
     }
-  }, [execute]);
+  }, changes);
 };