fix: wrong focus behavior when configuring menu selectability

src/Menu.tsx

@@ -387,24 +387,48 @@ const Menu = React.forwardRef<MenuRef, MenuProps>((props, ref) => {
     setMergedActiveKey(undefined);
   });
 
+  // ======================== Select ========================
+  // >>>>> Select keys
+  const [internalSelectKeys, setMergedSelectKeys] = useControlledState(
+    defaultSelectedKeys || [],
+    selectedKeys,
+  );
+  const mergedSelectKeys = React.useMemo(() => {
+    if (Array.isArray(internalSelectKeys)) {
+      return internalSelectKeys;
+    }
+
+    if (internalSelectKeys === null || internalSelectKeys === undefined) {
+      return EMPTY_LIST;
+    }
+
+    return [internalSelectKeys];
+  }, [internalSelectKeys]);
+  // >>>>> accept ref
   useImperativeHandle(ref, () => {
     return {
       list: containerRef.current,
       focus: options => {
+        if (!containerRef.current) {
+          return;
+        }
         const keys = getKeys();
         const { elements, key2element, element2key } = refreshElements(keys, uuid);
         const focusableElements = getFocusableElements(containerRef.current, elements);
-
-        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;
-        }
+        const focusableKeys = new Set(
+          focusableElements.map(el => element2key.get(el)).filter(Boolean),
+        );
+        const defaultFocusKey = focusableElements[0]
+          ? element2key.get(focusableElements[0])
+          : childList.find(node => !node.props.disabled)?.key;
+        const selectedFocusKey = mergedSelectKeys.find(k => focusableKeys.has(k));
+        const activeFocusKey =
+          mergedActiveKey && key2element.has(mergedActiveKey) ? mergedActiveKey : undefined;
+
+        const shouldFocusKey = selectable
+          ? (selectedFocusKey ?? activeFocusKey ?? defaultFocusKey)
+          : defaultFocusKey;
         const elementToFocus = key2element.get(shouldFocusKey);
-
         if (shouldFocusKey && elementToFocus) {
           elementToFocus?.focus?.(options);
         }
@@ -417,24 +441,6 @@ const Menu = React.forwardRef<MenuRef, MenuProps>((props, ref) => {
     };
   });
 
-  // ======================== Select ========================
-  // >>>>> Select keys
-  const [internalSelectKeys, setMergedSelectKeys] = useControlledState(
-    defaultSelectedKeys || [],
-    selectedKeys,
-  );
-  const mergedSelectKeys = React.useMemo(() => {
-    if (Array.isArray(internalSelectKeys)) {
-      return internalSelectKeys;
-    }
-
-    if (internalSelectKeys === null || internalSelectKeys === undefined) {
-      return EMPTY_LIST;
-    }
-
-    return [internalSelectKeys];
-  }, [internalSelectKeys]);
-
   // >>>>> Trigger select
   const triggerSelection = (info: MenuInfo) => {
     if (selectable) {

tests/Focus.spec.tsx

@@ -1,8 +1,10 @@
 /* eslint-disable no-undef */
 import { act, fireEvent, render } from '@testing-library/react';
 import { spyElementPrototypes } from '@rc-component/util/lib/test/domHook';
+import KeyCode from '@rc-component/util/lib/KeyCode';
 import React from 'react';
 import Menu, { MenuItem, MenuItemGroup, MenuRef, SubMenu } from '../src';
+import { isActive } from './util';
 
 describe('Focus', () => {
   beforeAll(() => {
@@ -26,6 +28,24 @@ describe('Focus', () => {
     jest.useRealTimers();
   });
 
+  function keyDown(container: HTMLElement, keyCode: number) {
+    fireEvent.keyDown(container.querySelector('ul.rc-menu-root'), {
+      which: keyCode,
+      keyCode,
+      charCode: keyCode,
+    });
+
+    // SubMenu raf need slow than accessibility
+    for (let i = 0; i < 20; i += 1) {
+      act(() => {
+        jest.advanceTimersByTime(10);
+      });
+    }
+    act(() => {
+      jest.runAllTimers();
+    });
+  }
+
   it('Get focus', async () => {
     const { container } = await act(async () =>
       render(
@@ -186,5 +206,129 @@ describe('Focus', () => {
     expect(document.activeElement).toBe(getByTitle('Submenu'));
     expect(getByTestId('sub-menu')).toHaveClass('rc-menu-submenu-active');
   });
+
+  it('When selectable is not configured, the focus should move to the first available item instead of keeping the previously focused item', async () => {
+    const menuRef = React.createRef<MenuRef>();
+    const items = [
+      { key: '0', label: 'First Item' },
+      { key: '1', label: 'Second Item' },
+      { key: '2', label: 'Third Item' },
+    ];
+    const TestApp = () => {
+      return (
+        <div>
+          <Menu data-testid="menu" selectable={false} ref={menuRef}>
+            {items.map(item => (
+              <MenuItem key={item.key} data-testid={item.key}>
+                {item.label}
+              </MenuItem>
+            ))}
+          </Menu>
+        </div>
+      );
+    };
+    const { getByTestId, container } = render(<TestApp />);
+    let focusSpy: jest.SpyInstance | undefined;
+    try {
+      // ================ check keydown ==============
+      // first item
+      keyDown(container, KeyCode.DOWN);
+      isActive(container, 0);
+      // second item
+      keyDown(container, KeyCode.DOWN);
+      isActive(container, 1);
+      // select second item
+      keyDown(container, KeyCode.ENTER);
+
+      // mock focus on item 0 to make sure it gets focused
+      const item0 = getByTestId('0');
+      focusSpy = jest.spyOn(item0, 'focus').mockImplementation(() => {});
+      menuRef.current.focus();
+      expect(focusSpy).toHaveBeenCalled();
+
+      // ================ check click ==============
+      // click third item
+      const item2 = getByTestId('2');
+      fireEvent.click(item2);
+      menuRef.current.focus();
+      expect(focusSpy).toHaveBeenCalled();
+    } finally {
+      focusSpy?.mockRestore();
+    }
+  });
+  it('When selectable is configured, the focus should move to the selected item if there is a selection, else to the first item, not retain on last focused item', async () => {
+    const menuRef = React.createRef<MenuRef>();
+    const items = [
+      { key: '0', label: 'First Item' },
+      { key: '1', label: 'Second Item' },
+      { key: '2', label: 'Third Item' },
+    ];
+    const TestApp = () => {
+      return (
+        <div>
+          <Menu data-testid="menu" selectable ref={menuRef}>
+            {items.map(item => (
+              <MenuItem key={item.key} data-testid={item.key}>
+                {item.label}
+              </MenuItem>
+            ))}
+          </Menu>
+        </div>
+      );
+    };
+    const { getByTestId, container } = render(<TestApp />);
+    let focusSpy: jest.SpyInstance | undefined;
+    let focusSpy2: jest.SpyInstance | undefined;
+    try {
+      // ================ check keydown ==============
+      // first item
+      keyDown(container, KeyCode.DOWN);
+      isActive(container, 0);
+      // second item
+      keyDown(container, KeyCode.DOWN);
+      isActive(container, 1);
+      // select second item
+      keyDown(container, KeyCode.ENTER);
+      // mock focus on item 1 to make sure it gets focused
+      const item1 = getByTestId('1');
+      focusSpy = jest.spyOn(item1, 'focus').mockImplementation(() => {});
+      menuRef.current.focus();
+      expect(focusSpy).toHaveBeenCalled();
+
+      // ================ check click ==============
+      // click third item
+      const item2 = getByTestId('2');
+      focusSpy2 = jest.spyOn(item2, 'focus').mockImplementation(() => {});
+      fireEvent.click(item2);
+      menuRef.current.focus();
+      // mock focus on item 2 to make sure it gets focused
+      expect(focusSpy2).toHaveBeenCalled();
+    } finally {
+      focusSpy?.mockRestore();
+      focusSpy2?.mockRestore();
+    }
+  });
+  it('should fallback when selected item is disabled', () => {
+    const menuRef = React.createRef<MenuRef>();
+    const items = [
+      { key: '1', label: 'Disabled', disabled: true },
+      { key: '2', label: 'Active' },
+      { key: '3', label: 'Item 3' },
+    ];
+    const { getByTestId } = render(
+      <Menu ref={menuRef} selectable selectedKeys={['disabled']}>
+        {items.map(item => (
+          <MenuItem disabled={item.disabled} key={item.key} data-testid={item.key}>
+            {item.label}
+          </MenuItem>
+        ))}
+      </Menu>,
+    );
+    const item2 = getByTestId('2');
+    const focusSpy = jest.spyOn(item2, 'focus').mockImplementation(() => {});
+    menuRef.current.focus();
+    expect(focusSpy).toHaveBeenCalled();
+    focusSpy?.mockRestore();
+  });
 });
 /* eslint-enable */