From f2265a72f8f37448e68420e194e8f98f36f7b429 Mon Sep 17 00:00:00 2001
From: Anton Korzunov <thekashey@gmail.com>
Date: Sat, 24 Aug 2024 22:15:07 +1000
Subject: [PATCH] fix: auto-focus in case of activeElement disappeareance,
 fixes #321

---
 .size-limit.json          |  6 +--
 _tests/keep-focus.spec.js | 92 ++++++++++++++++++++++++++++++++++++---
 src/Trap.js               | 30 ++++++++++++-
 3 files changed, 118 insertions(+), 10 deletions(-)

diff --git a/.size-limit.json b/.size-limit.json
index 705b8a8..e3b6b05 100644
--- a/.size-limit.json
+++ b/.size-limit.json
@@ -1,7 +1,7 @@
 [
   {
     "path": "dist/cjs/UI.js",
-    "limit": "3.8 KB",
+    "limit": "3.9 KB",
     "ignore": [
       "prop-types",
       "@babel/runtime",
@@ -10,7 +10,7 @@
   },
   {
     "path": "dist/es2015/sidecar.js",
-    "limit": "4.6 KB",
+    "limit": "4.8 KB",
     "ignore": [
       "prop-types",
       "@babel/runtime",
@@ -19,7 +19,7 @@
   },
   {
     "path": "dist/es2015/index.js",
-    "limit": "6.9 KB",
+    "limit": "7.1 KB",
     "ignore": [
       "prop-types",
       "@babel/runtime",
diff --git a/_tests/keep-focus.spec.js b/_tests/keep-focus.spec.js
index ee0a507..0b98d18 100644
--- a/_tests/keep-focus.spec.js
+++ b/_tests/keep-focus.spec.js
@@ -1,20 +1,102 @@
 import React from 'react';
 import {
-  render, screen,
+  render,
 } from '@testing-library/react';
 import { expect } from 'chai';
 import FocusLock from '../src';
+import { deferAction } from '../src/util';
 
-describe('Focus restoration', () => {
+describe('Focus restoration', async () => {
   it('maintains focus on element removal', async () => {
+    const { rerender } = render(
+      <FocusLock>
+        <button type="button" key={1}>btn1</button>
+      </FocusLock>,
+    );
+    //
+    expect(document.activeElement.innerHTML).to.be.equal('btn1');
+
+    rerender(
+      <FocusLock>
+        <button type="button" key={2}>new button</button>
+      </FocusLock>,
+    );
+
+    // wait
+    await new Promise(deferAction);
+
+    expect(document.activeElement.innerHTML).to.be.equal('new button');
+  });
+
+  it('handles disabled elements', async () => {
+    const { rerender } = render(
+      <FocusLock>
+        <button type="button">btn1</button>
+        <button type="button">btn2</button>
+      </FocusLock>,
+    );
+    //
+    expect(document.activeElement.innerHTML).to.be.equal('btn1');
+
+
+    // https://github.com/jsdom/jsdom/issues/3029 - jsdom does trigger blur on disabled
+    document.activeElement.blur();
+    document.body.focus();
+
+    rerender(
+      <FocusLock>
+        <button type="button" disabled>btn1</button>
+        <button type="button">btn2</button>
+      </FocusLock>,
+    );
+
+    // wait
+    await new Promise(deferAction);
+
+    expect(document.activeElement.innerHTML).to.be.equal('btn2');
+  });
+
+  it('moves focus to the nearest element', async () => {
     render(
       <FocusLock>
-        <button type="button">test</button>
+        <button type="button">btn1</button>
+        <button type="button">btn2</button>
+        <button type="button" id="middle-button">btn3</button>
+        <button type="button">btn4</button>
+        <button type="button">btn5</button>
       </FocusLock>,
     );
     //
-    expect(document.activeElement).to.be.equal(screen.getByRole('button'));
+    const middleButton = document.getElementById('middle-button');
+    middleButton.focus();
+    expect(document.activeElement).to.be.equal(middleButton);
+
+    middleButton.parentElement.removeChild(middleButton);
+    // wait
+    await new Promise(deferAction);
+
+    // btn4 "replaces" bnt3 in visual order
+    expect(document.activeElement.innerHTML).to.be.equal('btn4');
   });
 
-  it.todo('selects closes element to restore focus');
+  it('moves focus to the nearest element before', async () => {
+    render(
+      <FocusLock>
+        <button type="button">btn1</button>
+        <button type="button">btn2</button>
+        <button type="button" id="middle-button">btn3</button>
+      </FocusLock>,
+    );
+    //
+    const middleButton = document.getElementById('middle-button');
+    middleButton.focus();
+    expect(document.activeElement).to.be.equal(middleButton);
+
+    middleButton.parentElement.removeChild(middleButton);
+    // wait
+    await new Promise(deferAction);
+
+    // btn2 is just before bnt3
+    expect(document.activeElement.innerHTML).to.be.equal('btn2');
+  });
 });
diff --git a/src/Trap.js b/src/Trap.js
index e6bbd3a..489a10a 100644
--- a/src/Trap.js
+++ b/src/Trap.js
@@ -4,7 +4,9 @@ import PropTypes from 'prop-types';
 import withSideEffect from 'react-clientside-effect';
 import {
   moveFocusInside, focusInside,
-  focusIsHidden, expandFocusableNodes,
+  focusIsHidden,
+  expandFocusableNodes,
+  getFocusableNodes,
   focusNextElement,
   focusPrevElement,
   focusFirstElement,
@@ -22,6 +24,7 @@ const isFreeFocus = () => focusOnBody() || focusIsHidden();
 
 let lastActiveTrap = null;
 let lastActiveFocus = null;
+let tryRestoreFocus = () => null;
 
 let lastPortaledElement = null;
 
@@ -86,6 +89,10 @@ const withinHost = (activeElement, workingArea) => (
   workingArea.some(area => checkInHost(activeElement, area, area))
 );
 
+const isNotFocusable = node => (
+  !getFocusableNodes([node.parentNode], new Map()).some(el => el.node === node)
+);
+
 const activateTrap = () => {
   let result = false;
   if (lastActiveTrap) {
@@ -93,6 +100,24 @@ const activateTrap = () => {
       observed, persistentFocus, autoFocus, shards, crossFrame, focusOptions,
     } = lastActiveTrap;
     const workingNode = observed || (lastPortaledElement && lastPortaledElement.portaledElement);
+
+    // check if lastActiveFocus is still reachable
+    if (focusOnBody() && lastActiveFocus) {
+      if (
+        // it was removed
+        !document.body.contains(lastActiveFocus)
+          // or not focusable (this is expensive operation)!
+          || isNotFocusable(lastActiveFocus)
+      ) {
+        lastActiveFocus = null;
+
+        const newTarget = tryRestoreFocus();
+        if (newTarget) {
+          newTarget.focus();
+        }
+      }
+    }
+
     const activeElement = document && document.activeElement;
     if (workingNode) {
       const workingArea = [
@@ -130,11 +155,12 @@ const activateTrap = () => {
           }
           focusWasOutsideWindow = false;
           lastActiveFocus = document && document.activeElement;
+          tryRestoreFocus = captureFocusRestore(lastActiveFocus);
         }
       }
 
       if (document
-          // element was changed
+          // element was changed by moveFocusInside
           && activeElement !== document.activeElement
           // fast check for any auto-guard
           && document.querySelector('[data-focus-auto-guard]')) {