From d4e78c42a94be027b4dc7ed2659a5fddfbf9bd4e Mon Sep 17 00:00:00 2001
From: Jack Pope <jackpope1@gmail.com>
Date: Tue, 23 Apr 2024 12:13:09 -0400
Subject: [PATCH] Add ref callback test for cleanup fn vs null call (#28895)

Used this test scenario to clarify how callback refs work when detached
based on the availability of a cleanup function to update documentation
in https://github.com/reactjs/react.dev/pull/6770

Checking it in for additional test coverage and test-based documentation
---
 packages/react-dom/src/__tests__/refs-test.js | 97 +++++++++++++++++++
 1 file changed, 97 insertions(+)

diff --git a/packages/react-dom/src/__tests__/refs-test.js b/packages/react-dom/src/__tests__/refs-test.js
index 0c47e1728b..d68c6f9334 100644
--- a/packages/react-dom/src/__tests__/refs-test.js
+++ b/packages/react-dom/src/__tests__/refs-test.js
@@ -706,4 +706,101 @@ describe('refs return clean up function', () => {
     expect(setup).toHaveBeenCalledTimes(1);
     expect(cleanUp).toHaveBeenCalledTimes(1);
   });
+
+  it('handles detaching refs with either cleanup function or null argument', async () => {
+    const container = document.createElement('div');
+    const cleanUp = jest.fn();
+    const setup = jest.fn();
+    const setup2 = jest.fn();
+    const nullHandler = jest.fn();
+
+    function _onRefChangeWithCleanup(_ref) {
+      if (_ref) {
+        setup(_ref.id);
+      } else {
+        nullHandler();
+      }
+      return cleanUp;
+    }
+
+    function _onRefChangeWithoutCleanup(_ref) {
+      if (_ref) {
+        setup2(_ref.id);
+      } else {
+        nullHandler();
+      }
+    }
+
+    const root = ReactDOMClient.createRoot(container);
+    await act(() => {
+      root.render(<div id="test-div" ref={_onRefChangeWithCleanup} />);
+    });
+
+    expect(setup).toBeCalledWith('test-div');
+    expect(setup).toHaveBeenCalledTimes(1);
+    expect(cleanUp).toHaveBeenCalledTimes(0);
+
+    await act(() => {
+      root.render(<div id="test-div2" ref={_onRefChangeWithoutCleanup} />);
+    });
+
+    // Existing setup call was not called again
+    expect(setup).toHaveBeenCalledTimes(1);
+    // No null call because cleanup is returned
+    expect(nullHandler).toHaveBeenCalledTimes(0);
+    // Now we have a cleanup
+    expect(cleanUp).toHaveBeenCalledTimes(1);
+
+    // New ref is setup
+    expect(setup2).toBeCalledWith('test-div2');
+    expect(setup2).toHaveBeenCalledTimes(1);
+
+    // Now, render with the original ref again
+    await act(() => {
+      root.render(<div id="test-div3" ref={_onRefChangeWithCleanup} />);
+    });
+
+    // Setup was not called again
+    expect(setup2).toBeCalledWith('test-div2');
+    expect(setup2).toHaveBeenCalledTimes(1);
+
+    // Null handler hit because no cleanup is returned
+    expect(nullHandler).toHaveBeenCalledTimes(1);
+
+    // Original setup hit one more time
+    expect(setup).toHaveBeenCalledTimes(2);
+  });
+
+  it('calls cleanup function on unmount', async () => {
+    const container = document.createElement('div');
+    const cleanUp = jest.fn();
+    const setup = jest.fn();
+    const nullHandler = jest.fn();
+
+    function _onRefChangeWithCleanup(_ref) {
+      if (_ref) {
+        setup(_ref.id);
+      } else {
+        nullHandler();
+      }
+      return cleanUp;
+    }
+
+    const root = ReactDOMClient.createRoot(container);
+    await act(() => {
+      root.render(<div id="test-div" ref={_onRefChangeWithCleanup} />);
+    });
+
+    expect(setup).toHaveBeenCalledTimes(1);
+    expect(cleanUp).toHaveBeenCalledTimes(0);
+    expect(nullHandler).toHaveBeenCalledTimes(0);
+
+    root.unmount();
+
+    expect(setup).toHaveBeenCalledTimes(1);
+    // Now cleanup has been called
+    expect(cleanUp).toHaveBeenCalledTimes(1);
+    // Ref callback never called with null when cleanup is returned
+    expect(nullHandler).toHaveBeenCalledTimes(0);
+  });
 });