From 75284c62c8e4a68dfeb41a8d98a1e636e9ef531a Mon Sep 17 00:00:00 2001
From: xobotyi <xog3@yandex.ru>
Date: Fri, 22 Nov 2019 13:49:21 +0300
Subject: [PATCH] fix(#792): make useUnmount invoke the current callback
 version instead of very first

---
 src/useUnmount.ts        | 10 +++++--
 tests/useUnmount.test.ts | 57 ++++++++++++++++++++++++++--------------
 2 files changed, 46 insertions(+), 21 deletions(-)

diff --git a/src/useUnmount.ts b/src/useUnmount.ts
index c27bd7ea71..2bd460e959 100644
--- a/src/useUnmount.ts
+++ b/src/useUnmount.ts
@@ -1,7 +1,13 @@
+import { useRef } from 'react';
 import useEffectOnce from './useEffectOnce';
 
-const useUnmount = (fn: () => void | undefined) => {
-  useEffectOnce(() => fn);
+const useUnmount = (fn: () => any): void => {
+  const fnRef = useRef(fn);
+
+  // update the ref each render so if it change the newest callback will be invoked
+  fnRef.current = fn;
+
+  useEffectOnce(() => () => fnRef.current());
 };
 
 export default useUnmount;
diff --git a/tests/useUnmount.test.ts b/tests/useUnmount.test.ts
index 2c6078a99d..77fe892eff 100644
--- a/tests/useUnmount.test.ts
+++ b/tests/useUnmount.test.ts
@@ -1,32 +1,51 @@
 import { renderHook } from '@testing-library/react-hooks';
 import { useUnmount } from '../src';
 
-const mockCallback = jest.fn();
+describe('useUnmount', () => {
+  it('should be defined', () => {
+    expect(useUnmount).toBeDefined();
+  });
 
-afterEach(() => {
-  jest.resetAllMocks();
-});
+  it('should not call provided callback on mount', () => {
+    const spy = jest.fn();
+    renderHook(() => useUnmount(spy));
 
-it('should not call provided callback on mount', () => {
-  renderHook(() => useUnmount(mockCallback));
+    expect(spy).not.toHaveBeenCalled();
+  });
 
-  expect(mockCallback).not.toHaveBeenCalled();
-});
+  it('should not call provided callback on re-renders', () => {
+    const spy = jest.fn();
+    const hook = renderHook(() => useUnmount(spy));
 
-it('should call provided callback on unmount', () => {
-  const { unmount } = renderHook(() => useUnmount(mockCallback));
-  expect(mockCallback).not.toHaveBeenCalled();
+    hook.rerender();
+    hook.rerender();
+    hook.rerender();
+    hook.rerender();
 
-  unmount();
+    expect(spy).not.toHaveBeenCalled();
+  });
 
-  expect(mockCallback).toHaveBeenCalledTimes(1);
-});
+  it('should call provided callback on unmount', () => {
+    const spy = jest.fn();
+    const hook = renderHook(() => useUnmount(spy));
+
+    hook.unmount();
+
+    expect(spy).toHaveBeenCalledTimes(1);
+  });
 
-it('should not call provided callback on rerender', () => {
-  const { rerender } = renderHook(() => useUnmount(mockCallback));
-  expect(mockCallback).not.toHaveBeenCalled();
+  it('should call provided callback if is has been changed', () => {
+    const spy = jest.fn();
+    const spy2 = jest.fn();
+    const spy3 = jest.fn();
+    const hook = renderHook(cb => useUnmount(cb), { initialProps: spy });
 
-  rerender();
+    hook.rerender(spy2);
+    hook.rerender(spy3);
+    hook.unmount();
 
-  expect(mockCallback).not.toHaveBeenCalled();
+    expect(spy).not.toHaveBeenCalled();
+    expect(spy2).not.toHaveBeenCalled();
+    expect(spy3).toHaveBeenCalledTimes(1);
+  });
 });