From 1fb3294decef2bb9f16f164cee80c5f76127a0e1 Mon Sep 17 00:00:00 2001
From: Jack Pope <jackpope@meta.com>
Date: Tue, 23 Apr 2024 11:25:13 -0400
Subject: [PATCH 1/4] Add test for cleanup vs null call

---
 packages/react-dom/src/__tests__/refs-test.js | 66 +++++++++++++++++++
 1 file changed, 66 insertions(+)

diff --git a/packages/react-dom/src/__tests__/refs-test.js b/packages/react-dom/src/__tests__/refs-test.js
index 0c47e1728b270..357991b56ac81 100644
--- a/packages/react-dom/src/__tests__/refs-test.js
+++ b/packages/react-dom/src/__tests__/refs-test.js
@@ -706,4 +706,70 @@ 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 cleanUp2 = 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);
+    expect(cleanUp2).toHaveBeenCalledTimes(0);
+
+    // 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);
+  });
 });

From c484eb4363206c111de35d894cd0256e100e38a2 Mon Sep 17 00:00:00 2001
From: Jack Pope <jackpope@meta.com>
Date: Tue, 23 Apr 2024 11:39:13 -0400
Subject: [PATCH 2/4] f


From 8476bbfb40f018cc0d759501a685a8d61a0ca333 Mon Sep 17 00:00:00 2001
From: Jack Pope <jackpope@meta.com>
Date: Tue, 23 Apr 2024 11:42:54 -0400
Subject: [PATCH 3/4] remove unused function

---
 packages/react-dom/src/__tests__/refs-test.js | 2 --
 1 file changed, 2 deletions(-)

diff --git a/packages/react-dom/src/__tests__/refs-test.js b/packages/react-dom/src/__tests__/refs-test.js
index 357991b56ac81..c4b96a3fc808a 100644
--- a/packages/react-dom/src/__tests__/refs-test.js
+++ b/packages/react-dom/src/__tests__/refs-test.js
@@ -710,7 +710,6 @@ describe('refs return clean up function', () => {
   it('handles detaching refs with either cleanup function or null argument', async () => {
     const container = document.createElement('div');
     const cleanUp = jest.fn();
-    const cleanUp2 = jest.fn();
     const setup = jest.fn();
     const setup2 = jest.fn();
     const nullHandler = jest.fn();
@@ -755,7 +754,6 @@ describe('refs return clean up function', () => {
     // New ref is setup
     expect(setup2).toBeCalledWith('test-div2');
     expect(setup2).toHaveBeenCalledTimes(1);
-    expect(cleanUp2).toHaveBeenCalledTimes(0);
 
     // Now, render with the original ref again
     await act(() => {

From 2ab758faf3e810e9c25889b0f55f5604d3a05193 Mon Sep 17 00:00:00 2001
From: Jack Pope <jackpope@meta.com>
Date: Tue, 23 Apr 2024 11:57:42 -0400
Subject: [PATCH 4/4] Add additional test case to cover unmount behavior

---
 packages/react-dom/src/__tests__/refs-test.js | 33 +++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/packages/react-dom/src/__tests__/refs-test.js b/packages/react-dom/src/__tests__/refs-test.js
index c4b96a3fc808a..d68c6f93345e7 100644
--- a/packages/react-dom/src/__tests__/refs-test.js
+++ b/packages/react-dom/src/__tests__/refs-test.js
@@ -770,4 +770,37 @@ describe('refs return clean up function', () => {
     // 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);
+  });
 });