Fix focus restoring when activation happened lazily

packages/material-ui-unstyled/src/Unstable_TrapFocus/Unstable_TrapFocus.js

@@ -315,11 +315,23 @@ function Unstable_TrapFocus(props) {
     }
   };
 
+  const handleFocusSentinel = (event) => {
+    if (nodeToRestore.current === null) {
+      nodeToRestore.current = event.relatedTarget;
+    }
+    activated.current = true;
+  };
+
   return (
     <React.Fragment>
-      <div tabIndex={0} ref={sentinelStart} data-test="sentinelStart" />
+      <div
+        tabIndex={0}
+        onFocus={handleFocusSentinel}
+        ref={sentinelStart}
+        data-test="sentinelStart"
+      />
       {React.cloneElement(children, { ref: handleRef, onFocus })}
-      <div tabIndex={0} ref={sentinelEnd} data-test="sentinelEnd" />
+      <div tabIndex={0} onFocus={handleFocusSentinel} ref={sentinelEnd} data-test="sentinelEnd" />
     </React.Fragment>
   );
 }

test/e2e/fixtures/Unstable_TrapFocus/DefaultOpenLazyTrapFocus.tsx

@@ -0,0 +1,23 @@
+import * as React from 'react';
+import TrapFocus from '@material-ui/core/Unstable_TrapFocus';
+
+export default function BaseTrapFocus() {
+  const [open, close] = React.useReducer(() => false, true);
+
+  return (
+    <React.Fragment>
+      <button type="button" autoFocus data-testid="initial-focus">
+        initial focus
+      </button>
+      <TrapFocus isEnabled={() => true} open={open} disableAutoFocus>
+        <div data-testid="root">
+          <div>Title</div>
+          <button type="button" onClick={close}>
+            close
+          </button>
+          <button type="button">noop</button>
+        </div>
+      </TrapFocus>
+    </React.Fragment>
+  );
+}

test/e2e/index.test.ts

@@ -137,6 +137,21 @@ describe('e2e', () => {
       await expect(screen.getByText('ok')).toHaveFocus();
     });
 
+    it('should loop the tab key after activation', async () => {
+      await renderFixture('Unstable_TrapFocus/DefaultOpenLazyTrapFocus');
+
+      await expect(screen.getByTestId('initial-focus')).toHaveFocus();
+
+      await page.keyboard.press('Tab');
+      await expect(screen.getByText('close')).toHaveFocus();
+      await page.keyboard.press('Tab');
+      await expect(screen.getByText('noop')).toHaveFocus();
+      await page.keyboard.press('Tab');
+      await expect(screen.getByText('close')).toHaveFocus();
+      await page.keyboard.press('Enter');
+      await expect(screen.getByTestId('initial-focus')).toHaveFocus();
+    });
+
     it('should focus on first focus element after last has received a tab click', async () => {
       await renderFixture('Unstable_TrapFocus/OpenTrapFocus');