From 0ce421ced78fd6eb06a5676fefb856e18bfcacc1 Mon Sep 17 00:00:00 2001
From: Arnaud de Surirey <adesurirey@gmail.com>
Date: Tue, 7 Jan 2020 19:08:38 +0100
Subject: [PATCH] fix: useUpdateEffect returns optional cleanup function

---
 src/useUpdateEffect.ts        |  4 +++-
 tests/useUpdateEffect.test.ts | 21 +++++++++++++++------
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/src/useUpdateEffect.ts b/src/useUpdateEffect.ts
index c180d73700..736abf8c0c 100644
--- a/src/useUpdateEffect.ts
+++ b/src/useUpdateEffect.ts
@@ -5,7 +5,9 @@ const useUpdateEffect: typeof useEffect = (effect, deps) => {
   const isFirstMount = useFirstMountState();
 
   useEffect(() => {
-    !isFirstMount && effect();
+    if (!isFirstMount) {
+      return effect();
+    }
   }, deps);
 };
 
diff --git a/tests/useUpdateEffect.test.ts b/tests/useUpdateEffect.test.ts
index da55e36f03..5d00265bab 100644
--- a/tests/useUpdateEffect.test.ts
+++ b/tests/useUpdateEffect.test.ts
@@ -1,13 +1,22 @@
 import { renderHook } from '@testing-library/react-hooks';
 import { useUpdateEffect } from '../src';
 
-const mockEffectCleanup = jest.fn();
-const mockEffectCallback = jest.fn().mockReturnValue(mockEffectCleanup);
-
 it('should run effect on update', () => {
-  const { rerender } = renderHook(() => useUpdateEffect(mockEffectCallback));
-  expect(mockEffectCallback).not.toHaveBeenCalled();
+  const effect = jest.fn();
+
+  const { rerender } = renderHook(() => useUpdateEffect(effect));
+  expect(effect).not.toHaveBeenCalled();
 
   rerender();
-  expect(mockEffectCallback).toHaveBeenCalledTimes(1);
+  expect(effect).toHaveBeenCalledTimes(1);
+});
+
+it('should run cleanup on unmount', () => {
+  const cleanup = jest.fn();
+  const hook = renderHook(() => useUpdateEffect(cleanup));
+
+  hook.rerender();
+  hook.unmount();
+
+  expect(cleanup).toHaveBeenCalledTimes(1);
 });