feat(react-dom): whileElementsMounted callback (floating-ui#1683)

Author: atomiks

package-lock.json

@@ -33,18 +33,24 @@ export type UseFloatingReturn<RT extends ReferenceType = ReferenceType> =
     };
   };
 
-export interface Props {
+export interface Props<RT extends ReferenceType = ReferenceType> {
   open: boolean;
   onOpenChange: (open: boolean) => void;
   placement: Placement;
   middleware: Array<Middleware>;
   strategy: Strategy;
   nodeId: string;
+  whileElementsMounted?: (
+    reference: RT,
+    floating: HTMLElement,
+    update: () => void
+  ) => void | (() => void);
 }
 
 export function useFloating<RT extends ReferenceType = ReferenceType>({
   open = false,
   onOpenChange = () => {},
+  whileElementsMounted,
   placement,
   middleware,
   strategy,
@@ -55,7 +61,12 @@ export function useFloating<RT extends ReferenceType = ReferenceType>({
   const tree = useFloatingTree<RT>();
   const dataRef = React.useRef<ContextData>({});
   const events = React.useState(() => createPubSub())[0];
-  const floating = usePositionalFloating<RT>({placement, middleware, strategy});
+  const floating = usePositionalFloating<RT>({
+    placement,
+    middleware,
+    strategy,
+    whileElementsMounted,
+  });
 
   const context = React.useMemo<FloatingContext<RT>>(
     () => ({

packages/react-dom-interactions/src/index.ts

@@ -71,8 +71,7 @@
     "@rollup/plugin-commonjs": "^21.0.1",
     "@rollup/plugin-node-resolve": "^13.0.6",
     "@rollup/plugin-replace": "^3.0.0",
-    "@testing-library/react": "^12.1.2",
-    "@testing-library/react-hooks": "^7.0.2",
+    "@testing-library/react": "^13.2.0",
     "@types/jest": "^27.0.3",
     "@types/react": "^18.0.1",
     "babel-plugin-annotate-pure-calls": "^0.4.0",

packages/react-dom/package.json

@@ -31,16 +31,37 @@ export type UseFloatingReturn<RT extends ReferenceType = ReferenceType> =
     };
   };
 
+export type UseFloatingProps<RT extends ReferenceType = ReferenceType> = Omit<
+  Partial<ComputePositionConfig>,
+  'platform'
+> & {
+  whileElementsMounted?: (
+    reference: RT,
+    floating: HTMLElement,
+    update: () => void
+  ) => void | (() => void);
+};
+
+function useLatestRef<T>(value: T) {
+  const ref = useRef(value);
+  useLayoutEffect(() => {
+    ref.current = value;
+  });
+  return ref;
+}
+
 export function useFloating<RT extends ReferenceType = ReferenceType>({
   middleware,
   placement = 'bottom',
   strategy = 'absolute',
-}: Omit<
-  Partial<ComputePositionConfig>,
-  'platform'
-> = {}): UseFloatingReturn<RT> {
+  whileElementsMounted,
+}: UseFloatingProps = {}): UseFloatingReturn<RT> {
   const reference = useRef<RT | null>(null);
   const floating = useRef<HTMLElement | null>(null);
+
+  const whileElementsMountedRef = useLatestRef(whileElementsMounted);
+  const cleanupRef = useRef<void | (() => void) | null>(null);
+
   const [data, setData] = useState<Data>({
     // Setting these to `null` will allow the consumer to determine if
     // `computePosition()` has run yet
@@ -90,20 +111,41 @@ export function useFloating<RT extends ReferenceType = ReferenceType>({
 
   useLayoutEffect(update, [update]);
 
+  const runElementMountCallback = useCallback(() => {
+    if (typeof cleanupRef.current === 'function') {
+      cleanupRef.current();
+      cleanupRef.current = null;
+    }
+
+    if (
+      reference.current &&
+      floating.current &&
+      whileElementsMountedRef.current
+    ) {
+      cleanupRef.current = whileElementsMountedRef.current(
+        reference.current,
+        floating.current,
+        update
+      );
+    }
+  }, [update, whileElementsMountedRef]);
+
   const setReference: UseFloatingReturn<RT>['reference'] = useCallback(
     (node) => {
       reference.current = node;
+      runElementMountCallback();
       update();
     },
-    [update]
+    [update, runElementMountCallback]
   );
 
   const setFloating: UseFloatingReturn<RT>['floating'] = useCallback(
     (node) => {
       floating.current = node;
+      runElementMountCallback();
       update();
     },
-    [update]
+    [update, runElementMountCallback]
   );
 
   const refs = useMemo(() => ({reference, floating}), []);

packages/react-dom/src/index.ts

@@ -11,16 +11,8 @@ import {
   hide,
   limitShift,
 } from '../src';
-import {renderHook} from '@testing-library/react-hooks';
-import {render, waitFor, fireEvent} from '@testing-library/react';
-import {useRef, useState} from 'react';
-
-test('`x` and `y` are initially `null`', async () => {
-  const {result} = renderHook(() => useFloating());
-
-  expect(result.current.x).toBe(null);
-  expect(result.current.y).toBe(null);
-});
+import {render, fireEvent, screen, cleanup, act} from '@testing-library/react';
+import {useRef, useState, useEffect} from 'react';
 
 test('middleware is always fresh and does not cause an infinite loop', async () => {
   function InlineMiddleware() {
@@ -129,16 +121,142 @@ test('middleware is always fresh and does not cause an infinite loop', async ()
     );
   }
 
-  await waitFor(() => render(<InlineMiddleware />));
+  render(<InlineMiddleware />);
+
+  const {getByTestId} = render(<StateMiddleware />);
+  fireEvent.click(getByTestId('step1'));
+
+  await act(async () => {});
 
-  const {getByTestId} = await waitFor(() => render(<StateMiddleware />));
-  await waitFor(() => fireEvent.click(getByTestId('step1')));
-  await waitFor(() => expect(getByTestId('x').textContent).toBe('10'));
+  expect(getByTestId('x').textContent).toBe('10');
 
-  await waitFor(() => fireEvent.click(getByTestId('step2')));
-  await waitFor(() => expect(getByTestId('x').textContent).toBe('5'));
+  fireEvent.click(getByTestId('step2'));
+
+  await act(async () => {});
+
+  expect(getByTestId('x').textContent).toBe('5');
 
   // No `expect` as this test will fail if a render loop occurs
-  await waitFor(() => fireEvent.click(getByTestId('step3')));
-  await waitFor(() => fireEvent.click(getByTestId('step4')));
+  fireEvent.click(getByTestId('step3'));
+  fireEvent.click(getByTestId('step4'));
+
+  await act(async () => {});
+});
+
+describe('whileElementsMounted', () => {
+  test('is called a single time when both elements mount', () => {
+    const spy = jest.fn();
+
+    function App() {
+      const {reference, floating} = useFloating({whileElementsMounted: spy});
+      return (
+        <>
+          <button ref={reference} />
+          <div ref={floating} />
+        </>
+      );
+    }
+
+    render(<App />);
+    expect(spy).toHaveBeenCalledTimes(1);
+    cleanup();
+  });
+
+  test('is called a single time after floating mounts conditionally', () => {
+    const spy = jest.fn();
+
+    function App() {
+      const [open, setOpen] = useState(false);
+      const {reference, floating} = useFloating({whileElementsMounted: spy});
+      return (
+        <>
+          <button ref={reference} onClick={() => setOpen(true)} />
+          {open && <div ref={floating} />}
+        </>
+      );
+    }
+
+    render(<App />);
+    expect(spy).toHaveBeenCalledTimes(0);
+    fireEvent.click(screen.getByRole('button'));
+    expect(spy).toHaveBeenCalledTimes(1);
+
+    cleanup();
+  });
+
+  test('is called a single time after reference mounts conditionally', () => {
+    const spy = jest.fn();
+
+    function App() {
+      const [open, setOpen] = useState(false);
+      const {reference, floating} = useFloating({whileElementsMounted: spy});
+      return (
+        <>
+          {open && <button ref={reference} />}
+          <div role="tooltip" ref={floating} onClick={() => setOpen(true)} />
+        </>
+      );
+    }
+
+    render(<App />);
+    expect(spy).toHaveBeenCalledTimes(0);
+    fireEvent.click(screen.getByRole('tooltip'));
+    expect(spy).toHaveBeenCalledTimes(1);
+
+    cleanup();
+  });
+
+  test('is called a single time both elements mount conditionally', () => {
+    const spy = jest.fn();
+
+    function App() {
+      const [open, setOpen] = useState(false);
+      const {reference, floating} = useFloating({whileElementsMounted: spy});
+
+      useEffect(() => {
+        setOpen(true);
+      }, []);
+
+      return (
+        <>
+          {open && <button ref={reference} />}
+          {open && <div role="tooltip" ref={floating} />}
+        </>
+      );
+    }
+
+    render(<App />);
+    expect(spy).toHaveBeenCalledTimes(1);
+
+    cleanup();
+  });
+
+  test('calls the cleanup function', () => {
+    const cleanupSpy = jest.fn();
+    const spy = jest.fn(() => cleanupSpy);
+
+    function App() {
+      const [open, setOpen] = useState(true);
+      const {reference, floating} = useFloating({whileElementsMounted: spy});
+
+      useEffect(() => {
+        setOpen(false);
+      }, []);
+
+      return (
+        <>
+          {open && <button ref={reference} />}
+          {open && <div role="tooltip" ref={floating} />}
+        </>
+      );
+    }
+
+    render(<App />);
+    expect(cleanupSpy).toHaveBeenCalledTimes(1);
+
+    // Does not get called again post-cleanup
+    expect(spy).toHaveBeenCalledTimes(1);
+
+    cleanup();
+  });
 });