fix: support callback forwarded refs

src/useEnsuredForwardedRef.ts

@@ -1,4 +1,5 @@
 import {
+  ForwardedRef,
   forwardRef,
   ForwardRefExoticComponent,
   MutableRefObject,
@@ -10,16 +11,29 @@ import {
   useRef,
 } from 'react';
 
+const assignRef = <T>(ref: ForwardedRef<T>, value: T | null) => {
+  if (!ref) {
+    return;
+  }
+
+  if (typeof ref === 'function') {
+    ref(value);
+  } else {
+    ref.current = value;
+  }
+};
+
 export default function useEnsuredForwardedRef<T>(
-  forwardedRef: MutableRefObject<T>
+  forwardedRef: ForwardedRef<T>
 ): MutableRefObject<T> {
-  const ensuredRef = useRef(forwardedRef && forwardedRef.current);
+  const ensuredRef = useRef<T>(
+    (forwardedRef && typeof forwardedRef !== 'function' ? forwardedRef.current : null) as T
+  );
 
   useEffect(() => {
-    if (!forwardedRef) {
-      return;
-    }
-    forwardedRef.current = ensuredRef.current;
+    assignRef(forwardedRef, ensuredRef.current);
+
+    return () => assignRef(forwardedRef, null);
   }, [forwardedRef]);
 
   return ensuredRef;
@@ -29,7 +43,7 @@ export function ensuredForwardRef<T, P = {}>(
   Component: RefForwardingComponent<T, P>
 ): ForwardRefExoticComponent<PropsWithoutRef<P> & RefAttributes<T>> {
   return forwardRef((props: PropsWithChildren<P>, ref) => {
-    const ensuredRef = useEnsuredForwardedRef(ref as MutableRefObject<T>);
+    const ensuredRef = useEnsuredForwardedRef(ref);
     return Component(props, ensuredRef);
   });
 }

tests/useEnsuredForwardedRef.test.tsx

@@ -72,3 +72,24 @@ test('should return a valid ref when using the wrapper function style', () => {
   expect(initialRef.current).toBeTruthy();
   expect(initialRef.current?.id).toBe('test_id');
 });
+
+test('should call callback refs when using the wrapper function style', () => {
+  const callbackRef = jest.fn();
+
+  const WrappedComponent = ensuredForwardRef<HTMLDivElement>((_props, ref) => {
+    return <div id="test_id" ref={ref} />;
+  });
+
+  TestUtils.act(() => {
+    ReactDOM.render(<WrappedComponent ref={callbackRef} />, container);
+  });
+
+  expect(callbackRef).toHaveBeenLastCalledWith(expect.any(HTMLDivElement));
+  expect(callbackRef.mock.calls[0][0].id).toBe('test_id');
+
+  TestUtils.act(() => {
+    ReactDOM.unmountComponentAtNode(container);
+  });
+
+  expect(callbackRef).toHaveBeenLastCalledWith(null);
+});