From 8215cb96af64aba1cf84b31e86d480668bf8a37e Mon Sep 17 00:00:00 2001
From: Stalgia Grigg <stalgia@bocoup.com>
Date: Tue, 12 Dec 2023 07:59:44 -0800
Subject: [PATCH] Trap backwards navigation from initial focus on header in
 focus-trapped modals (#851)

* Add unit test for reverse navigation from initial focused header, FocusTrapper

* Update unit tests for FocusTrapper, handle backwards navigation from initially focused header

* Handle focus leaving non-focusable header element backwards in FocusTrapper
---
 client/components/common/BasicModal/index.jsx |  8 +++++-
 .../components/common/FocusTrapper/index.jsx  | 27 ++++++++++++++-----
 client/tests/FocusTrapper.test.jsx            | 19 ++++++++++---
 3 files changed, 42 insertions(+), 12 deletions(-)

diff --git a/client/components/common/BasicModal/index.jsx b/client/components/common/BasicModal/index.jsx
index 936a85b89..8b4d52dda 100644
--- a/client/components/common/BasicModal/index.jsx
+++ b/client/components/common/BasicModal/index.jsx
@@ -48,7 +48,13 @@ const BasicModal = ({
     }, []);
 
     return (
-        <FocusTrapper isActive={show} trappedElId={`focus-trapped-${id}`}>
+        <FocusTrapper
+            isActive={show}
+            initialFocusRef={
+                initialFocusRef?.current ? initialFocusRef : headerRef
+            }
+            trappedElId={`focus-trapped-${id}`}
+        >
             <Modal
                 show={show}
                 id={`focus-trapped-${id}`}
diff --git a/client/components/common/FocusTrapper/index.jsx b/client/components/common/FocusTrapper/index.jsx
index 77406fdf5..60e27642f 100644
--- a/client/components/common/FocusTrapper/index.jsx
+++ b/client/components/common/FocusTrapper/index.jsx
@@ -2,9 +2,8 @@ import React from 'react';
 import PropTypes from 'prop-types';
 import { useEffect, useRef } from 'react';
 
-const FocusTrapper = ({ children, isActive, trappedElId }) => {
+const FocusTrapper = ({ children, isActive, initialFocusRef, trappedElId }) => {
     const focusableElsRef = useRef([]);
-
     const updateFocusableElements = () => {
         if (focusableElsRef.current.length > 0) return;
 
@@ -41,13 +40,22 @@ const FocusTrapper = ({ children, isActive, trappedElId }) => {
         const firstFocusableEl = focusableEls[1];
         // Last focusable element is before the blank 'after' div
         const lastFocusableEl = focusableEls[focusableEls.length - 2];
-
-        if (event.shiftKey && document.activeElement === firstFocusableEl) {
+        // When SHIFT + TAB is pressed and the active element is the first focusable element
+        if (
+            event.shiftKey &&
+            (document.activeElement === firstFocusableEl ||
+                document.activeElement === focusableEls[0] ||
+                document.activeElement === initialFocusRef.current)
+        ) {
             lastFocusableEl.focus();
             event.preventDefault();
-        } else if (
+        }
+        // When TAB is pressed and the active element is the last focusable element
+        else if (
             !event.shiftKey &&
-            document.activeElement === lastFocusableEl
+            (document.activeElement === lastFocusableEl ||
+                document.activeElement ===
+                    focusableEls[focusableEls.length - 1])
         ) {
             firstFocusableEl.focus();
             event.preventDefault();
@@ -60,9 +68,11 @@ const FocusTrapper = ({ children, isActive, trappedElId }) => {
             document.addEventListener('keydown', trapFocus);
         } else {
             document.removeEventListener('keydown', trapFocus);
+            focusableElsRef.current = [];
         }
 
         return () => {
+            focusableElsRef.current = [];
             document.removeEventListener('keydown', trapFocus);
         };
     }, [isActive]);
@@ -73,7 +83,10 @@ const FocusTrapper = ({ children, isActive, trappedElId }) => {
 FocusTrapper.propTypes = {
     children: PropTypes.node.isRequired,
     isActive: PropTypes.bool.isRequired,
-    trappedElId: PropTypes.string.isRequired
+    trappedElId: PropTypes.string.isRequired,
+    initialFocusRef: PropTypes.shape({
+        current: PropTypes.object
+    })
 };
 
 export default FocusTrapper;
diff --git a/client/tests/FocusTrapper.test.jsx b/client/tests/FocusTrapper.test.jsx
index 7144e38bf..a4c114d85 100644
--- a/client/tests/FocusTrapper.test.jsx
+++ b/client/tests/FocusTrapper.test.jsx
@@ -6,7 +6,7 @@ import { render, fireEvent, act } from '@testing-library/react';
 import FocusTrapper from '../components/common/FocusTrapper';
 
 describe('FocusTrapper', () => {
-    let trappedDiv;
+    let trappedDiv, initialFocusRef;
 
     beforeEach(() => {
         trappedDiv = document.createElement('div');
@@ -19,8 +19,20 @@ describe('FocusTrapper', () => {
     });
 
     const renderEls = async () => {
+        initialFocusRef = React.createRef();
         return render(
-            <FocusTrapper isActive={true} trappedElId="trapped-div">
+            <FocusTrapper
+                isActive={true}
+                initialFocusRef={initialFocusRef}
+                trappedElId="trapped-div"
+            >
+                <h1
+                    className="modal-header"
+                    tabIndex="-1"
+                    ref={initialFocusRef}
+                >
+                    Modal Header
+                </h1>
                 <button>Click Me</button>
                 <a href="www">Link</a>
             </FocusTrapper>,
@@ -67,11 +79,10 @@ describe('FocusTrapper', () => {
 
         const container = document.getElementById('trapped-div');
 
-        const firstFocusable = container.querySelector('button');
         const lastFocusable = container.querySelector('a');
 
         act(() => {
-            firstFocusable.focus();
+            initialFocusRef.current.focus();
         });
 
         act(() => {