fix: focusTriggerAfterClose not working

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,56 @@ describe('dialog', () => {
     expect(modalRender.find('.rc-dialog-content').props().style.background).toEqual('#1890ff');
   });
 
+  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 />);
+    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());
+  });
+
+  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 />);
+    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());
+  });
+
   describe('size should work', () => {
     it('width', () => {
       const wrapper = mount(<Dialog visible width={1128} />);