fix(Menu): primitive portal via Theme refs (#1100)

src/core/primitives/Menu/fragments/MenuPrimitivePortal.tsx

@@ -2,16 +2,20 @@
 import React, { useContext, useEffect, useState, forwardRef, ElementRef, ComponentPropsWithoutRef } from 'react';
 import Floater from '~/core/primitives/Floater';
 import MenuPrimitiveRootContext from '../contexts/MenuPrimitiveRootContext';
+import ThemeContext from '~/components/ui/Theme/ThemeContext';
 
 export type MenuPrimitivePortalElement = HTMLDivElement;
 export type MenuPrimitivePortalProps = { children: React.ReactNode } & ComponentPropsWithoutRef<typeof Floater.Portal>;
 
 const MenuPrimitivePortal = forwardRef<MenuPrimitivePortalElement, MenuPrimitivePortalProps>(
     ({ children, ...props }, ref) => {
         const context = useContext(MenuPrimitiveRootContext);
+        const themeContext = useContext(ThemeContext);
         const [rootElementFound, setRootElementFound] = useState(false);
         const rootElement = (
-            document.querySelector('#rad-ui-theme-container') || document.body
+            themeContext?.portalRootRef.current
+            ?? themeContext?.containerRef.current
+            ?? document.body
         ) as HTMLElement | null;
 
         useEffect(() => {

src/core/primitives/Menu/tests/MenuPrimitive.portal.test.tsx

@@ -0,0 +1,46 @@
+import React from 'react';
+import { render, screen, waitFor } from '@testing-library/react';
+import userEvent from '@testing-library/user-event';
+import Theme from '~/components/ui/Theme/Theme';
+import DropdownMenu from '~/components/ui/DropdownMenu/DropdownMenu';
+
+const mockMatchMedia = () => {
+    if ('matchMedia' in window && typeof window.matchMedia === 'function') return;
+    Object.defineProperty(window, 'matchMedia', {
+        writable: true,
+        value: jest.fn().mockImplementation(() => ({
+            matches: false,
+            addEventListener: jest.fn(),
+            removeEventListener: jest.fn()
+        }))
+    });
+};
+
+describe('Menu primitive portal', () => {
+    beforeEach(() => mockMatchMedia());
+
+    test('portals menu into Theme portal root', async() => {
+        const user = userEvent.setup();
+
+        render(
+            <Theme>
+                <DropdownMenu.Root>
+                    <DropdownMenu.Trigger>Open</DropdownMenu.Trigger>
+                    <DropdownMenu.Portal>
+                        <DropdownMenu.Content>
+                            <DropdownMenu.Item label="Profile">Profile</DropdownMenu.Item>
+                        </DropdownMenu.Content>
+                    </DropdownMenu.Portal>
+                </DropdownMenu.Root>
+            </Theme>
+        );
+
+        const portalRoot = document.querySelector('[data-rad-ui-portal-root]') as HTMLElement;
+        expect(portalRoot).toBeTruthy();
+
+        await user.click(screen.getByText('Open'));
+        await waitFor(() => {
+            expect(portalRoot).toContainElement(screen.getByText('Profile'));
+        });
+    });
+});

src/test-utils/portal.ts

@@ -1,17 +1,37 @@
+import React from 'react';
 import { render } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
+import ThemeContext from '~/components/ui/Theme/ThemeContext';
 
 export function renderWithPortal(ui: React.ReactElement) {
     const portalRoot = document.createElement('div');
-    portalRoot.setAttribute('id', 'rad-ui-theme-container');
+    portalRoot.setAttribute('data-rad-ui-portal-root', '');
     document.body.appendChild(portalRoot);
-    const result = render(ui);
+
+    const themeContainer = document.createElement('div');
+    themeContainer.setAttribute('id', 'rad-ui-theme-container');
+    document.body.appendChild(themeContainer);
+
+    const result = render(
+        React.createElement(
+            ThemeContext.Provider,
+            {
+                value: {
+                    containerRef: { current: themeContainer },
+                    portalRootRef: { current: portalRoot }
+                }
+            },
+            ui
+        )
+    );
+
     return {
         ...result,
         portalRoot,
         cleanup: () => {
             result.unmount();
             portalRoot.remove();
+            themeContainer.remove();
         }
     };
 }
@@ -28,7 +48,6 @@ export async function assertFocusTrap(container: HTMLElement) {
     }
     const first = focusable[0];
     const last = focusable[focusable.length - 1];
-    // TODO: track jsdom focus trap flakiness and re-enable assertions in CI when stable.
     const isJsdom = typeof navigator !== 'undefined' && /jsdom/i.test(navigator.userAgent);
 
     first.focus();