fix: focusTriggerAfterClose not working (#307)

close #280
ant-design/ant-design#32923

Author: afc163

src/Dialog/index.tsx

@@ -51,14 +51,23 @@ export default function Dialog(props: IDialogPropTypes) {
   // ========================== Init ==========================
   const ariaId = useId();
 
+  function saveLastOutSideActiveElementRef() {
+    if (!contains(wrapperRef.current, document.activeElement)) {
+      lastOutSideActiveElementRef.current = document.activeElement as HTMLElement;
+    }
+  }
+
+  function focusDialogContent() {
+    if (!contains(wrapperRef.current, document.activeElement)) {
+      contentRef.current?.focus();
+    }
+  }
+
   // ========================= Events =========================
   function onDialogVisibleChanged(newVisible: boolean) {
+    // Try to focus
     if (newVisible) {
-      // Try to focus
-      if (!contains(wrapperRef.current, document.activeElement)) {
-        lastOutSideActiveElementRef.current = document.activeElement as HTMLElement;
-        contentRef.current?.focus();
-      }
+      focusDialogContent();
     } else {
       // Clean up scroll bar & focus back
       setAnimatedVisible(false);
@@ -131,6 +140,7 @@ export default function Dialog(props: IDialogPropTypes) {
   useEffect(() => {
     if (visible) {
       setAnimatedVisible(true);
+      saveLastOutSideActiveElementRef();
     }
   }, [visible]);
 

tests/index.spec.tsx

@@ -1,5 +1,5 @@
 /* eslint-disable react/no-render-return-value, max-classes-per-file, func-names, no-console */
-import React, { cloneElement } from 'react';
+import React, { cloneElement, useEffect } from 'react';
 import { act } from 'react-dom/test-utils';
 import { render } from '@testing-library/react';
 import type { ReactWrapper } from 'enzyme';
@@ -362,6 +362,60 @@ describe('dialog', () => {
     expect(modalRender.find('.rc-dialog-content').props().style.background).toEqual('#1890ff');
   });
 
+  describe('focusTriggerAfterClose', () => {
+    it('should focus trigger after close dialog', () => {
+      const Demo = () => {
+        const [visible, setVisible] = React.useState(false);
+        return (
+          <>
+            <button onClick={() => setVisible(true)}>trigger</button>
+            <Dialog visible={visible} onClose={() => setVisible(false)}>
+              content
+            </Dialog>
+          </>
+        );
+      };
+      const wrapper = mount(<Demo />, { attachTo: document.body });
+      const trigger = wrapper.find('button').at(0);
+      (trigger.getDOMNode() as any).focus();
+      trigger.simulate('click');
+      jest.runAllTimers();
+      const closeButton = wrapper.find('.rc-dialog-close');
+      closeButton.simulate('click');
+      jest.runAllTimers();
+      expect(document.activeElement).toBe(trigger.getDOMNode());
+      wrapper.unmount();
+    });
+
+    it('should focus trigger after close dialog when contains focusable element', () => {
+      const Demo = () => {
+        const [visible, setVisible] = React.useState(false);
+        const inputRef = React.useRef(null);
+        useEffect(() => {
+          inputRef.current?.focus();
+        }, []);
+        return (
+          <>
+            <button onClick={() => setVisible(true)}>trigger</button>
+            <Dialog visible={visible} onClose={() => setVisible(false)}>
+              <input ref={inputRef} />
+            </Dialog>
+          </>
+        );
+      };
+      const wrapper = mount(<Demo />, { attachTo: document.body });
+      const trigger = wrapper.find('button').at(0);
+      (trigger.getDOMNode() as any).focus();
+      trigger.simulate('click');
+      jest.runAllTimers();
+      const closeButton = wrapper.find('.rc-dialog-close');
+      closeButton.simulate('click');
+      jest.runAllTimers();
+      expect(document.activeElement).toBe(trigger.getDOMNode());
+      wrapper.unmount();
+    });
+  });
+
   describe('size should work', () => {
     it('width', () => {
       const wrapper = mount(<Dialog visible width={1128} />);