refactor(fuselage-hooks): useSafeRefCallback reimplemented callback/cleanup order

.changeset/chatty-ways-tan.md

@@ -0,0 +1,5 @@
+---
+"@rocket.chat/fuselage-hooks": patch
+---
+
+refactor(fuselage-hooks): `useSafeRefCallback` reimplemented callback/cleanup order

packages/fuselage-hooks/src/useSafeRefCallback/useSafeRefCallback.spec.tsx

@@ -28,14 +28,12 @@ describe('useSafeRefCallback', () => {
 
     rerender(<TestComponent callback={callback} renderSpan />);
 
-    expect(callback).toHaveBeenCalledTimes(3);
-    expect(callback.mock.calls[1][0]).toBe(null);
-    expect(callback.mock.calls[2][0]).toBeInstanceOf(HTMLSpanElement);
+    expect(callback).toHaveBeenCalledTimes(2);
+    expect(callback.mock.lastCall[0]).toBeInstanceOf(HTMLSpanElement);
 
     unmount();
 
-    expect(callback).toHaveBeenCalledTimes(4);
-    expect(callback.mock.calls[3][0]).toBe(null);
+    expect(callback).toHaveBeenCalledTimes(2);
   });
 
   it('should run again when callback reference changes', () => {
@@ -50,23 +48,58 @@ describe('useSafeRefCallback', () => {
 
     rerender(<TestComponent callback={callback2} />);
 
+    expect(callback).toHaveBeenCalledTimes(1);
+
+    expect(callback2).toHaveBeenCalledTimes(1);
+    expect(callback2.mock.lastCall[0]).toBeInstanceOf(HTMLDivElement);
+
+    rerender(<TestComponent callback={callback2} renderSpan />);
+
+    expect(callback2).toHaveBeenCalledTimes(2);
+    expect(callback2.mock.lastCall[0]).toBeInstanceOf(HTMLSpanElement);
+
+    unmount();
+
+    expect(callback2).toHaveBeenCalledTimes(2);
+  });
+
+  it('should call cleanup if callback changes', () => {
+    const cleanup = jest.fn();
+    const callback = jest.fn<() => void, any>(() => cleanup);
+
+    const { rerender, unmount } = render(<TestComponent callback={callback} />);
+
+    expect(callback).toHaveBeenCalledTimes(1);
+    expect(callback.mock.lastCall[0]).toBeInstanceOf(HTMLDivElement);
+    expect(cleanup).not.toHaveBeenCalled();
+
+    const cleanup2 = jest.fn();
+    const callback2 = jest.fn<() => void, any>(() => cleanup2);
+
+    rerender(<TestComponent callback={callback2} />);
+
     // Ensure first callback has been properly unmounted
-    expect(callback).toHaveBeenCalledTimes(2);
-    expect(callback.mock.calls[1][0]).toBe(null);
+    expect(callback).toHaveBeenCalledTimes(1);
+    expect(callback.mock.lastCall[0]).toBeInstanceOf(HTMLDivElement);
+
+    expect(cleanup).toHaveBeenCalledTimes(1);
 
     expect(callback2).toHaveBeenCalledTimes(1);
     expect(callback2.mock.lastCall[0]).toBeInstanceOf(HTMLDivElement);
+    expect(cleanup2).not.toHaveBeenCalled();
 
     rerender(<TestComponent callback={callback2} renderSpan />);
 
-    expect(callback2).toHaveBeenCalledTimes(3);
-    expect(callback2.mock.calls[1][0]).toBe(null);
-    expect(callback2.mock.calls[2][0]).toBeInstanceOf(HTMLSpanElement);
+    expect(callback2).toHaveBeenCalledTimes(2);
+    expect(callback2.mock.lastCall[0]).toBeInstanceOf(HTMLSpanElement);
+
+    expect(cleanup2).toHaveBeenCalledTimes(1);
 
     unmount();
 
-    expect(callback2).toHaveBeenCalledTimes(4);
-    expect(callback2.mock.calls[3][0]).toBe(null);
+    expect(callback2).toHaveBeenCalledTimes(2);
+
+    expect(cleanup2).toHaveBeenCalledTimes(2);
   });
 
   it('should call cleanup with previous value on rerender', () => {
@@ -82,23 +115,18 @@ describe('useSafeRefCallback', () => {
 
     rerender(<TestComponent callback={callback} renderSpan />);
 
-    expect(callback).toHaveBeenCalledTimes(3);
-    expect(callback.mock.calls[1][0]).toBe(null);
-    expect(callback.mock.calls[2][0]).toBeInstanceOf(HTMLSpanElement);
+    expect(callback).toHaveBeenCalledTimes(2);
+    expect(callback.mock.lastCall[0]).toBeInstanceOf(HTMLSpanElement);
 
-    expect(cleanup).toHaveBeenCalledTimes(2);
+    expect(cleanup).toHaveBeenCalledTimes(1);
 
     const cleanup2 = jest.fn();
     const callback2 = jest.fn<() => void, any>(() => cleanup2);
 
     rerender(<TestComponent callback={callback2} renderSpan />);
 
-    // Ensure first callback has been properly unmounted
-    expect(callback).toHaveBeenCalledTimes(4);
-    expect(callback.mock.calls[3][0]).toBe(null);
-
-    console.log(cleanup.mock.calls);
-    expect(cleanup).toHaveBeenCalledTimes(3);
+    expect(callback).toHaveBeenCalledTimes(2);
+    expect(cleanup).toHaveBeenCalledTimes(2);
 
     expect(callback2).toHaveBeenCalledTimes(1);
     expect(callback2.mock.lastCall[0]).toBeInstanceOf(HTMLSpanElement);
@@ -107,17 +135,14 @@ describe('useSafeRefCallback', () => {
 
     rerender(<TestComponent callback={callback2} />);
 
-    expect(callback2).toHaveBeenCalledTimes(3);
-    expect(callback2.mock.calls[1][0]).toBe(null);
-    expect(callback2.mock.calls[2][0]).toBeInstanceOf(HTMLDivElement);
+    expect(callback2).toHaveBeenCalledTimes(2);
+    expect(callback2.mock.lastCall[0]).toBeInstanceOf(HTMLDivElement);
 
-    expect(cleanup2).toHaveBeenCalledTimes(2);
+    expect(cleanup2).toHaveBeenCalledTimes(1);
 
     unmount();
 
-    expect(callback2).toHaveBeenCalledTimes(4);
-    expect(callback2.mock.calls[3][0]).toBe(null);
-
-    expect(cleanup2).toHaveBeenCalledTimes(3);
+    expect(callback2).toHaveBeenCalledTimes(2);
+    expect(cleanup2).toHaveBeenCalledTimes(2);
   });
 });

packages/fuselage-hooks/src/useSafeRefCallback/useSafeRefCallback.ts

@@ -1,21 +1,19 @@
 import { useMemo } from 'react';
 
-type CallbackRefWithCleanup<T> = (node: T) => () => void;
-type CallbackRef<T> = (node: T) => void;
-
-type SafeCallbackRef<T> = CallbackRefWithCleanup<T> | CallbackRef<T>;
+type SafeCallbackRef<T> = (node: T) => (() => void) | undefined;
 
 /**
  * useSafeRefCallback will call a cleanup function (returned from the passed callback)
- * if the passed callback is called multiple times (similar to useEffect, but in a callbackRef)
+ * whenever the component is re-rendered ( similar to useEffect, but in a callbackRef )
+ *
+ * Caveat: Usually, callback refs are called with `null` when the component is unmounted. The returned callback from this hook will only be called whenever `node !== null`.
+ *  For cases where you want to do some action when the node is null, you should rely on the cleanup function to be called.
  *
  * @example
  *   const callback = useSafeRefCallback(
  *       useCallback(
- *           (node: T) => {
- *               if (!node) {
- *                   return;
- *               }
+ *            // node will always exist
+ *           (node: HTMLDivElement) => {
  *               node.addEventListener('click', listener);
  *               return () => {
  *                   node.removeEventListener('click', listener);
@@ -26,21 +24,21 @@ type SafeCallbackRef<T> = CallbackRefWithCleanup<T> | CallbackRef<T>;
  *   );
  *
  */
-export const useSafeRefCallback = <T extends HTMLElement | null>(
+export const useSafeRefCallback = <T extends HTMLElement>(
   callback: SafeCallbackRef<T>,
-) => {
-  const callbackRef = useMemo(() => {
-    let _cleanup: (() => void) | null;
+) =>
+  useMemo(() => {
+    let cleanup: (() => void) | null = null;
 
-    return (node: T): void => {
-      if (typeof _cleanup === 'function') {
-        _cleanup();
+    return (node: T | null): void => {
+      if (node === null) {
+        if (typeof cleanup === 'function') {
+          cleanup();
+          cleanup = null;
+        }
+        return;
       }
-      const cleanup = callback(node);
 
-      _cleanup = cleanup || null;
+      cleanup = callback(node) || null;
     };
   }, [callback]);
-
-  return callbackRef;
-};