Fix focus scope restore losing focus (#2166)

* Fix focus scope restore losing focus

* Add a test for this case

* reseting focus to nodeToRestore if nodeToRestore is within a parent FocusScope

Co-authored-by: Daniel Lu <[email protected]>

Author: snowystinger

packages/@react-aria/focus/src/FocusScope.tsx

@@ -376,8 +376,14 @@ function useRestoreFocus(scopeRef: RefObject<HTMLElement[]>, restoreFocus: boole
         if (nextElement) {
           focusElement(nextElement, true);
         } else {
-          // If there is no next element, blur the focused element to move focus to the body.
-          focusedElement.blur();
+           // If there is no next element and the nodeToRestore isn't within a FocusScope (i.e. we are leaving the top level focus scope)
+           // then move focus to the body.
+           // Otherwise restore focus to the nodeToRestore (e.g menu within a popover -> tabbing to close the menu should move focus to menu trigger)
+          if (!isElementInAnyScope(nodeToRestore, scopes)) {
+            focusedElement.blur();
+          } else {
+            focusElement(nodeToRestore, true);
+          }
         }
       }
     };

packages/@react-spectrum/dialog/test/DialogTrigger.test.js

@@ -15,6 +15,7 @@ import {ActionButton, Button} from '@react-spectrum/button';
 import {ButtonGroup} from '@react-spectrum/buttongroup';
 import {Content} from '@react-spectrum/view';
 import {Dialog, DialogTrigger} from '../';
+import {Heading} from '@react-spectrum/text';
 import {Item, Menu, MenuTrigger} from '@react-spectrum/menu';
 import MatchMediaMock from 'jest-matchmedia-mock';
 import {Provider} from '@react-spectrum/provider';
@@ -28,9 +29,11 @@ import userEvent from '@testing-library/user-event';
 describe('DialogTrigger', function () {
   let matchMedia;
   beforeAll(() => {
+    jest.spyOn(window.screen, 'width', 'get').mockImplementation(() => 1024);
     jest.useFakeTimers();
   });
   afterAll(() => {
+    jest.clearAllMocks();
     jest.useRealTimers();
   });
 
@@ -927,4 +930,54 @@ describe('DialogTrigger', function () {
 
     expect(document.activeElement).toBe(innerInput);
   });
+
+  it('will not lose focus to body', async () => {
+    let {getByRole, getByTestId} = render(
+      <Provider theme={theme}>
+        <DialogTrigger type="popover">
+          <ActionButton>Trigger</ActionButton>
+          <Dialog>
+            <Heading>The Heading</Heading>
+            <Content>
+              <MenuTrigger>
+                <ActionButton data-testid="innerButton">Test</ActionButton>
+                <Menu autoFocus="first">
+                  <Item>Item 1</Item>
+                  <Item>Item 2</Item>
+                  <Item>Item 3</Item>
+                </Menu>
+              </MenuTrigger>
+            </Content>
+          </Dialog>
+        </DialogTrigger>
+      </Provider>
+    );
+    let button = getByRole('button');
+    triggerPress(button);
+
+    act(() => {
+      jest.runAllTimers();
+    });
+
+    let outerDialog = getByRole('dialog');
+
+    await waitFor(() => {
+      expect(outerDialog).toBeVisible();
+    }); // wait for animation
+    let innerButton = getByTestId('innerButton');
+    userEvent.tab();
+    fireEvent.keyDown(document.activeElement, {key: 'Enter'});
+    fireEvent.keyUp(document.activeElement, {key: 'Enter'});
+
+    act(() => {
+      jest.runAllTimers();
+    });
+    userEvent.tab();
+    act(() => {
+      jest.runAllTimers();
+    });
+
+    expect(document.activeElement).toBe(innerButton);
+  });
+
 });