fix: useUpdateEffect returns optional cleanup function (#864)

fix: useUpdateEffect returns optional cleanup function

src/useUpdateEffect.ts

@@ -5,7 +5,9 @@ const useUpdateEffect: typeof useEffect = (effect, deps) => {
   const isFirstMount = useFirstMountState();
 
   useEffect(() => {
-    !isFirstMount && effect();
+    if (!isFirstMount) {
+      return effect();
+    }
   }, deps);
 };
 

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);
 });