fix: wrong focus behavior when configuring menu selectability

src/Menu.tsx

@@ -394,14 +394,23 @@ const Menu = React.forwardRef<MenuRef, MenuProps>((props, ref) => {
         const keys = getKeys();
         const { elements, key2element, element2key } = refreshElements(keys, uuid);
         const focusableElements = getFocusableElements(containerRef.current, elements);
-
+        const defaultFocusKey = focusableElements[0]
+          ? element2key.get(focusableElements[0])
+          : childList.find(node => !node.props.disabled)?.key;
         let shouldFocusKey: string;
-        if (mergedActiveKey && keys.includes(mergedActiveKey)) {
-          shouldFocusKey = mergedActiveKey;
+        // find the item to focus on based on whether it is selectable
+        if (selectable) {
+          const mergedSelectKeys = getMergedSelectKeys();
+          // if there is already selected items, select first item to focus
+          if (mergedSelectKeys.length && keys.includes(mergedSelectKeys[0])) {
+            shouldFocusKey = mergedSelectKeys[0];
+          } else if (mergedActiveKey && keys.includes(mergedActiveKey)) {
+            shouldFocusKey = mergedActiveKey;
+          } else {
+            shouldFocusKey = defaultFocusKey;
+          }
         } else {
-          shouldFocusKey = focusableElements[0]
-            ? element2key.get(focusableElements[0])
-            : childList.find(node => !node.props.disabled)?.key;
+          shouldFocusKey = defaultFocusKey;
         }
         const elementToFocus = key2element.get(shouldFocusKey);
 
@@ -434,6 +443,9 @@ const Menu = React.forwardRef<MenuRef, MenuProps>((props, ref) => {
 
     return [internalSelectKeys];
   }, [internalSelectKeys]);
+  function getMergedSelectKeys() {
+    return mergedSelectKeys;
+  }
 
   // >>>>> Trigger select
   const triggerSelection = (info: MenuInfo) => {

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,100 @@ 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 />);
+    // ================ 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');
+    const 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();
+    // cleanup
+    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 />);
+    // ================ 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');
+    const focusSpy = jest.spyOn(item1, 'focus').mockImplementation(() => {});
+    menuRef.current.focus();
+    expect(focusSpy).toHaveBeenCalled();
+
+    // ================ check click ==============
+    // click third item
+    const item2 = getByTestId('2');
+    const 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();
+    // cleanup
+    focusSpy.mockRestore();
+    focusSpy2.mockRestore();
+  });
 });
 /* eslint-enable */