Fixed ref passing and added ref tests for BodyPortal

Dantemss · Dantemss · commit 8db17f57c593 · 2024-08-26T11:45:43.000-05:00
diff --git a/src/components/BodyPortal.spec.tsx b/src/components/BodyPortal.spec.tsx
@@ -1,3 +1,4 @@
+import React from 'react';
 import { render } from '@testing-library/react';
 import { BodyPortal } from './BodyPortal';
 import { BodyPortalSlotsContext } from './BodyPortalSlotsContext';
@@ -176,6 +177,61 @@ describe('BodyPortal', () => {
     Modal
   </div>
 </body>
+`);
+  });
+
+  it('accepts an optional ref parameter that will be set', () => {
+    const TestPortal = ({ children }: React.PropsWithChildren<{}>) => {
+      const ref = React.useRef<HTMLElement | null>(null);
+      expect(ref.current).toBeNull();
+
+      React.useEffect(() => {
+        expect(ref.current).toBeInstanceOf(HTMLElement);
+      }, []);
+
+      return <BodyPortal ref={ref} slot='footer' tagName='footer'>{children}</BodyPortal>;
+    };
+    render(<><TestPortal>Footer stuff</TestPortal><h1>Title</h1></>, { container: root });
+    expect(document.body).toMatchInlineSnapshot(`
+<body>
+  <main
+    id="root"
+  >
+    <h1>
+      Title
+    </h1>
+  </main>
+  <footer
+    data-portal-slot="footer"
+  >
+    Footer stuff
+  </footer>
+</body>
+`);
+  });
+
+  it('accepts a ref callback', () => {
+    const setRef = jest.fn().mockImplementation((element) => {
+      expect(element).toBeInstanceOf(HTMLElement);
+    });
+    render(<><BodyPortal ref={setRef} slot='footer' tagName='footer'>Footer stuff</BodyPortal>
+            <h1>Title</h1></>, { container: root });
+    expect(setRef).toHaveBeenCalled();
+    expect(document.body).toMatchInlineSnapshot(`
+<body>
+  <main
+    id="root"
+  >
+    <h1>
+      Title
+    </h1>
+  </main>
+  <footer
+    data-portal-slot="footer"
+  >
+    Footer stuff
+  </footer>
+</body>
 `);
   });
 });
diff --git a/src/components/BodyPortal.tsx b/src/components/BodyPortal.tsx
@@ -21,21 +21,24 @@ const getInsertBeforeTarget = (bodyPortalSlots: string[], slot?: string) => {
   return null;
 }
 
-export const BodyPortal = ({
-  children, className, ref, role, slot, tagName
-}: React.PropsWithChildren<{
+export const BodyPortal = React.forwardRef<HTMLElement, React.PropsWithChildren<{
   className?: string;
-  ref?: React.MutableRefObject<HTMLElement | null>;
   role?: string;
   slot?: string;
-  tagName?: string
-}>) => {
+  tagName?: string;
+}>>(({ children, className, role, slot, tagName }, ref) => {
   const tag = tagName?.toUpperCase() ?? 'DIV';
   const internalRef = React.useRef<HTMLElement>(document.createElement(tag));
   if (internalRef.current.tagName !== tag) {
     internalRef.current = document.createElement(tag);
   }
-  if (ref) { ref.current = internalRef.current; }
+  if (ref) {
+    if (typeof ref === 'function') {
+      ref(internalRef.current);
+    } else {
+      ref.current = internalRef.current;
+    }
+  }
 
   const bodyPortalOrderedRefs = React.useContext(BodyPortalSlotsContext);
 
@@ -64,4 +67,4 @@ export const BodyPortal = ({
   }, [bodyPortalOrderedRefs, className, role, slot, tag]);
 
   return createPortal(children, internalRef.current);
-};
+});
