fix: fix focus when active element is removed (#791)

Author: jon-cullison

src/Menu.tsx

@@ -399,12 +399,14 @@ const Menu = React.forwardRef<MenuRef, MenuProps>((props, ref) => {
         const { elements, key2element, element2key } = refreshElements(keys, uuid);
         const focusableElements = getFocusableElements(containerRef.current, elements);
 
-        const shouldFocusKey =
-          mergedActiveKey ??
-          (focusableElements[0]
+        let shouldFocusKey: string;
+        if (mergedActiveKey && keys.includes(mergedActiveKey)) {
+          shouldFocusKey = mergedActiveKey;
+        } else {
+          shouldFocusKey = focusableElements[0]
             ? element2key.get(focusableElements[0])
-            : childList.find(node => !node.props.disabled)?.key);
-
+            : childList.find(node => !node.props.disabled)?.key;
+        }
         const elementToFocus = key2element.get(shouldFocusKey);
 
         if (shouldFocusKey && elementToFocus) {

tests/Menu.spec.tsx

@@ -8,6 +8,7 @@ import { act } from 'react-dom/test-utils';
 import type { MenuRef } from '../src';
 import Menu, { Divider, MenuItem, MenuItemGroup, SubMenu } from '../src';
 import { isActive, last } from './util';
+import { spyElementPrototypes } from '@rc-component/util/lib/test/domHook';
 
 jest.mock('@rc-component/trigger', () => {
   const react = require('react');
@@ -790,5 +791,73 @@ describe('Menu', () => {
     expect(catItem).toBe(container.querySelectorAll('li')[1]);
     expect(nonExistentItem).toBe(null);
   });
+
+  it('should correctly handle invalid mergedActiveKey and fallback to focusable elements', async () => {
+    // Mock to force make menu item visible
+    const domSpy = spyElementPrototypes(HTMLElement, {
+      offsetParent: {
+        get() {
+          return this.parentElement;
+        },
+      },
+    });
+
+    const menuRef = React.createRef<MenuRef>();
+
+    const TestApp = () => {
+      const [items, setItems] = React.useState([
+        { key: 'item1', label: 'First Item' },
+        { key: 'item2', label: 'Second Item' },
+        { key: 'item3', label: 'Third Item' },
+      ]);
+
+      const removeSecondItem = () => {
+        setItems(prevItems => prevItems.filter(item => item.key !== 'item2'));
+      };
+
+      return (
+        <div>
+          <Menu data-testid="menu" activeKey="item2" tabIndex={0} ref={menuRef}>
+            {items.map(item => (
+              <MenuItem key={item.key} data-testid={item.key}>
+                {item.label}
+              </MenuItem>
+            ))}
+          </Menu>
+          <button data-testid="remove-button" onClick={removeSecondItem} tabIndex={0}>
+            Remove Second Item
+          </button>
+        </div>
+      );
+    };
+
+    const { getByTestId } = render(<TestApp />);
+
+    const removeButton = getByTestId('remove-button');
+
+    // Verify initial state - item2 should be active
+    expect(getByTestId('item2')).toHaveClass('rc-menu-item-active');
+
+    // Remove the active item (item2)
+    await act(async () => {
+      fireEvent.click(removeButton);
+    });
+
+    // Verify the item is removed
+    expect(() => getByTestId('item2')).toThrow();
+
+    // mock focus on item 1 to make sure it gets focused
+    const item1 = getByTestId('item1');
+    const focusSpy = jest.spyOn(item1, 'focus').mockImplementation(() => {});
+
+    // when we call focus(), it should properly handle the case where
+    // mergedActiveKey ("item2") no longer exists
+    menuRef.current.focus();
+
+    expect(focusSpy).toHaveBeenCalled();
+    // cleanup
+    focusSpy.mockRestore();
+    domSpy?.mockRestore?.();
+  });
 });
 /* eslint-enable */