fix: math editor saves on close and focuses when switching node [PIE-695][PIE-697]

packages/editable-html-tip-tap/src/extensions/__tests__/math.test.js

@@ -648,6 +648,31 @@ describe('MathNodeView', () => {
     });
   });
 
+  it('re-registers click listener when node changes', async () => {
+    const addEventListenerSpy = jest.spyOn(document, 'addEventListener');
+    const removeEventListenerSpy = jest.spyOn(document, 'removeEventListener');
+    const nodeA = { attrs: { latex: 'x^2' } };
+    const nodeB = { attrs: { latex: 'y^2' } };
+
+    const { rerender } = render(<MathNodeView {...defaultProps} node={nodeA} selected={true} />);
+
+    await waitFor(() => {
+      expect(addEventListenerSpy).toHaveBeenCalledWith('click', expect.any(Function));
+    });
+
+    const initialCallCount = addEventListenerSpy.mock.calls.length;
+
+    rerender(<MathNodeView {...defaultProps} node={nodeB} selected={true} />);
+
+    await waitFor(() => {
+      expect(removeEventListenerSpy).toHaveBeenCalled();
+      expect(addEventListenerSpy.mock.calls.length).toBeGreaterThan(initialCallCount);
+    });
+
+    addEventListenerSpy.mockRestore();
+    removeEventListenerSpy.mockRestore();
+  });
+
   it('does not close toolbar when clicking the math node preview', async () => {
     const { getByTestId, queryByTestId } = render(<MathNodeView {...defaultProps} selected={true} />);
 

packages/editable-html-tip-tap/src/extensions/math.js

@@ -298,7 +298,7 @@ export const MathNodeView = (props) => {
     }
 
     return () => document.removeEventListener('click', handleClickOutside);
-  }, [editor, showToolbar]);
+  }, [editor, showToolbar, node]);
 
   return (
     <NodeViewWrapper

packages/math-toolbar/src/__tests__/editor-and-pad.test.js

@@ -9,8 +9,66 @@ describe('EditorAndPad', () => {
     onBlur: jest.fn(),
   };
 
+  beforeEach(() => {
+    jest.useFakeTimers();
+  });
+
+  afterEach(() => {
+    jest.useRealTimers();
+  });
+
   it('renders with default props', () => {
     const { container } = render(<EditorAndPad {...defaultProps} />);
     expect(container.firstChild).toBeInTheDocument();
   });
+
+  describe('autoFocus', () => {
+    it('focuses input immediately via setTimeout when autoFocus is true', () => {
+      const mockFocus = jest.fn();
+      const component = new EditorAndPad({ ...defaultProps, autoFocus: true });
+      component.input = { focus: mockFocus };
+
+      component.componentDidMount();
+
+      expect(mockFocus).not.toHaveBeenCalled();
+
+      jest.runAllTimers();
+
+      expect(mockFocus).toHaveBeenCalledTimes(1);
+    });
+
+    it('does not focus input when autoFocus is false', () => {
+      const mockFocus = jest.fn();
+      const component = new EditorAndPad({ ...defaultProps, autoFocus: false });
+      component.input = { focus: mockFocus };
+
+      component.componentDidMount();
+      jest.runAllTimers();
+
+      expect(mockFocus).not.toHaveBeenCalled();
+    });
+
+    it('does not focus when input ref is not set', () => {
+      const component = new EditorAndPad({ ...defaultProps, autoFocus: true });
+      component.input = null;
+
+      expect(() => {
+        component.componentDidMount();
+        jest.runAllTimers();
+      }).not.toThrow();
+    });
+
+    it('uses setTimeout with 0ms delay to defer focus', () => {
+      const setTimeoutSpy = jest.spyOn(global, 'setTimeout');
+      const mockFocus = jest.fn();
+      const component = new EditorAndPad({ ...defaultProps, autoFocus: true });
+      component.input = { focus: mockFocus };
+
+      component.componentDidMount();
+
+      expect(setTimeoutSpy).toHaveBeenCalledWith(expect.any(Function), 0);
+
+      setTimeoutSpy.mockRestore();
+    });
+  });
 });

packages/math-toolbar/src/editor-and-pad.jsx

@@ -263,7 +263,8 @@ export class EditorAndPad extends React.Component {
 
   componentDidMount() {
     if (this.input && this.props.autoFocus) {
-      this.input.focus();
+      // adding a timeout to wait for other stuff related to focus to be finished
+      setTimeout(() => this.input.focus(), 0);
     }
   }