From 3180f86371e249f366aec04fbb7839e609141a66 Mon Sep 17 00:00:00 2001
From: Laura Aubin <laura.aubin@shopify.com>
Date: Tue, 5 Dec 2023 14:57:02 -0500
Subject: [PATCH] [Modal] Fix activator not gaining focus when modal is closed
 (#11266)

Fix modal activator focus on close and add new optional activatorWrapper prop
---
 .changeset/empty-poets-wait.md                |  6 ++
 polaris-react/src/components/Box/Box.tsx      |  2 +-
 polaris-react/src/components/Modal/Modal.tsx  | 18 ++++-
 .../src/components/Modal/tests/Modal.test.tsx | 75 +++++++++++++++----
 4 files changed, 82 insertions(+), 19 deletions(-)
 create mode 100644 .changeset/empty-poets-wait.md

diff --git a/.changeset/empty-poets-wait.md b/.changeset/empty-poets-wait.md
new file mode 100644
index 00000000000..23dec26da68
--- /dev/null
+++ b/.changeset/empty-poets-wait.md
@@ -0,0 +1,6 @@
+---
+'@shopify/polaris': minor
+---
+
+- Fixed `Modal` `activator` not regaining focus on close
+- Added an `activatorWrapper` prop to `Modal` to support setting the element that wraps the `activator`
diff --git a/polaris-react/src/components/Box/Box.tsx b/polaris-react/src/components/Box/Box.tsx
index 4571a9335e1..ded4436ccdc 100644
--- a/polaris-react/src/components/Box/Box.tsx
+++ b/polaris-react/src/components/Box/Box.tsx
@@ -18,7 +18,7 @@ import type {ResponsiveProp} from '../../utilities/css';
 
 import styles from './Box.scss';
 
-type Element = 'div' | 'span' | 'section' | 'legend' | 'ul' | 'li';
+export type Element = 'div' | 'span' | 'section' | 'legend' | 'ul' | 'li';
 
 type LineStyles = 'solid' | 'dashed';
 type Overflow = 'hidden' | 'scroll';
diff --git a/polaris-react/src/components/Modal/Modal.tsx b/polaris-react/src/components/Modal/Modal.tsx
index 59345615f1d..d1aa8cb5b67 100644
--- a/polaris-react/src/components/Modal/Modal.tsx
+++ b/polaris-react/src/components/Modal/Modal.tsx
@@ -1,4 +1,4 @@
-import React, {useState, useCallback, useRef, useId, cloneElement} from 'react';
+import React, {useState, useCallback, useRef, useId} from 'react';
 import {TransitionGroup} from 'react-transition-group';
 
 import {focusFirstFocusableNode} from '../../utilities/focus';
@@ -7,6 +7,7 @@ import {WithinContentContext} from '../../utilities/within-content-context';
 import {wrapWithComponent} from '../../utilities/components';
 import {Backdrop} from '../Backdrop';
 import {Box} from '../Box';
+import type {Element} from '../Box';
 import {InlineStack} from '../InlineStack';
 import {Scrollable} from '../Scrollable';
 import {Spinner} from '../Spinner';
@@ -59,6 +60,11 @@ export interface ModalProps extends FooterProps {
   onScrolledToBottom?(): void;
   /** The element or the RefObject that activates the Modal */
   activator?: React.RefObject<HTMLElement> | React.ReactElement;
+  /**
+   * The element type to wrap the activator in
+   * @default 'div'
+   */
+  activatorWrapper?: Element;
   /** Removes Scrollable container from the modal content */
   noScroll?: boolean;
 }
@@ -82,6 +88,7 @@ export const Modal: React.FunctionComponent<ModalProps> & {
   secondaryActions,
   onScrolledToBottom,
   activator,
+  activatorWrapper = 'div',
   onClose,
   onIFrameLoad,
   onTransitionEnd,
@@ -112,6 +119,7 @@ export const Modal: React.FunctionComponent<ModalProps> & {
       activator && isRef(activator)
         ? activator && activator.current
         : activatorRef.current;
+
     if (activatorElement) {
       requestAnimationFrame(() => focusFirstFocusableNode(activatorElement));
     }
@@ -217,9 +225,11 @@ export const Modal: React.FunctionComponent<ModalProps> & {
   const animated = !instant;
 
   const activatorMarkup =
-    activator && !isRef(activator)
-      ? cloneElement(activator, {ref: activatorRef})
-      : null;
+    activator && !isRef(activator) ? (
+      <Box ref={activatorRef} as={activatorWrapper}>
+        {activator}
+      </Box>
+    ) : null;
 
   return (
     <WithinContentContext.Provider value>
diff --git a/polaris-react/src/components/Modal/tests/Modal.test.tsx b/polaris-react/src/components/Modal/tests/Modal.test.tsx
index 857e3021e16..a1a46c3b835 100644
--- a/polaris-react/src/components/Modal/tests/Modal.test.tsx
+++ b/polaris-react/src/components/Modal/tests/Modal.test.tsx
@@ -1,3 +1,4 @@
+import type {ReactNode} from 'react';
 import React, {useRef} from 'react';
 import {animationFrame} from '@shopify/jest-dom-mocks';
 import {mountWithApp} from 'tests/utilities';
@@ -27,17 +28,31 @@ jest.mock('react-transition-group', () => {
   };
 });
 
+jest.mock('../../TrapFocus', () => ({
+  ...jest.requireActual('../../TrapFocus'),
+  TrapFocus({children}: {children: ReactNode}) {
+    return <div>{children}</div>;
+  },
+}));
+
 describe('<Modal>', () => {
   let scrollSpy: jest.SpyInstance;
+  let requestAnimationFrameSpy: jest.SpyInstance;
 
   beforeEach(() => {
     scrollSpy = jest.spyOn(window, 'scroll');
     animationFrame.mock();
+    requestAnimationFrameSpy = jest.spyOn(window, 'requestAnimationFrame');
+    requestAnimationFrameSpy.mockImplementation((cb: () => number) => {
+      cb();
+      return 1;
+    });
   });
 
   afterEach(() => {
     animationFrame.restore();
     scrollSpy.mockRestore();
+    requestAnimationFrameSpy.mockRestore();
   });
 
   it('has a child with contentContext', () => {
@@ -513,10 +528,41 @@ describe('<Modal>', () => {
       }).not.toThrow();
     });
 
-    // Causes a circular dependency that causes the whole test file to be unrunnable
-    // eslint-disable-next-line jest/no-disabled-tests
-    it.skip('focuses the activator when the activator is an element on close', () => {
+    it('wraps the activator in a div by default', () => {
+      const activator = <Button />;
+
+      const modal = mountWithApp(
+        <Modal title="foo" onClose={noop} open={false} activator={activator} />,
+      );
+
+      expect(modal).toContainReactComponent(Box, {
+        as: 'div',
+        children: activator,
+      });
+    });
+
+    it('wraps the activator in a span if activatorWrapper is provided', () => {
+      const activator = <Button />;
+
+      const modal = mountWithApp(
+        <Modal
+          title="foo"
+          onClose={noop}
+          open={false}
+          activator={activator}
+          activatorWrapper="span"
+        />,
+      );
+
+      expect(modal).toContainReactComponent(Box, {
+        as: 'span',
+        children: activator,
+      });
+    });
+
+    it('focuses the activator on close when the activator is an element', () => {
       const id = 'activator-id';
+
       const modal = mountWithApp(
         <Modal
           title="foo"
@@ -527,21 +573,22 @@ describe('<Modal>', () => {
       );
 
       modal.find(Dialog)!.trigger('onExited');
-      const activator = modal.find('button', {id})!.domNode;
+
+      const activator = modal.find('button', {
+        id,
+      })!.domNode;
 
       expect(document.activeElement).toBe(activator);
     });
 
-    // Causes a circular dependency that causes the whole test file to be unrunnable
-    // eslint-disable-next-line jest/no-disabled-tests
-    it.skip('focuses the activator when the activator a ref on close', () => {
-      const buttonId = 'buttonId';
+    it('focuses the activator on close when the activator is a ref', () => {
+      const id = 'buttonId';
       const TestHarness = () => {
         const buttonRef = useRef<HTMLDivElement>(null);
 
         const button = (
           <div ref={buttonRef}>
-            <Button id={buttonId} />
+            <Button id={id} />
           </div>
         );
 
@@ -557,11 +604,11 @@ describe('<Modal>', () => {
 
       testHarness.find(Modal)!.find(Dialog)!.trigger('onExited');
 
-      expect(document.activeElement).toBe(
-        testHarness.findWhere(
-          (wrap) => wrap.is('button') && wrap.prop('id') === buttonId,
-        )!.domNode,
-      );
+      const activator = testHarness.find('button', {
+        id,
+      })!.domNode;
+
+      expect(document.activeElement).toBe(activator);
     });
   });
 });