From 6499db9c3cbbfba6d13688f1ee158e0e7208a8ae Mon Sep 17 00:00:00 2001
From: Sebastian Silbermann <silbermann.sebastian@gmail.com>
Date: Sun, 28 Jan 2024 10:41:47 +0100
Subject: [PATCH] Convert refs to createRoot (#28113)

---
 packages/react-dom/src/__tests__/refs-test.js | 357 +++++++++++-------
 .../src/__tests__/refsLegacy-test.js          | 101 +++++
 2 files changed, 327 insertions(+), 131 deletions(-)
 create mode 100644 packages/react-dom/src/__tests__/refsLegacy-test.js

diff --git a/packages/react-dom/src/__tests__/refs-test.js b/packages/react-dom/src/__tests__/refs-test.js
index 355ce3d18c93e..e42661abd5a1f 100644
--- a/packages/react-dom/src/__tests__/refs-test.js
+++ b/packages/react-dom/src/__tests__/refs-test.js
@@ -10,9 +10,10 @@
 'use strict';
 
 let React = require('react');
-let ReactDOM = require('react-dom');
+let ReactDOMClient = require('react-dom/client');
 let ReactFeatureFlags = require('shared/ReactFeatureFlags');
 let ReactTestUtils = require('react-dom/test-utils');
+let act = require('internal-test-utils').act;
 
 // This is testing if string refs are deleted from `instance.refs`
 // Once support for string refs is removed, this test can be removed.
@@ -23,9 +24,10 @@ describe('reactiverefs', () => {
   beforeEach(() => {
     jest.resetModules();
     React = require('react');
-    ReactDOM = require('react-dom');
+    ReactDOMClient = require('react-dom/client');
     ReactFeatureFlags = require('shared/ReactFeatureFlags');
     ReactTestUtils = require('react-dom/test-utils');
+    act = require('internal-test-utils').act;
   });
 
   afterEach(() => {
@@ -82,7 +84,7 @@ describe('reactiverefs', () => {
   /**
    * Render a TestRefsComponent and ensure that the main refs are wired up.
    */
-  const renderTestRefsComponent = function () {
+  const renderTestRefsComponent = async function () {
     /**
      * Only purpose is to test that refs are tracked even when applied to a
      * component that is injected down several layers. Ref systems are difficult to
@@ -121,8 +123,17 @@ describe('reactiverefs', () => {
     document.body.appendChild(container);
 
     let testRefsComponent;
-    expect(() => {
-      testRefsComponent = ReactDOM.render(<TestRefsComponent />, container);
+    await expect(async () => {
+      const root = ReactDOMClient.createRoot(container);
+      await act(() => {
+        root.render(
+          <TestRefsComponent
+            ref={current => {
+              testRefsComponent = current;
+            }}
+          />,
+        );
+      });
     }).toErrorDev([
       'Warning: Component "div" contains the string ref "resetDiv". ' +
         'Support for string refs will be removed in a future major release. ' +
@@ -157,8 +168,8 @@ describe('reactiverefs', () => {
    * Ensure that for every click log there is a corresponding ref (from the
    * perspective of the injected ClickCounter component.
    */
-  it('Should increase refs with an increase in divs', () => {
-    const testRefsComponent = renderTestRefsComponent();
+  it('Should increase refs with an increase in divs', async () => {
+    const testRefsComponent = await renderTestRefsComponent();
     const clickIncrementer = ReactTestUtils.findRenderedDOMComponentWithClass(
       testRefsComponent,
       'clickIncrementer',
@@ -171,14 +182,20 @@ describe('reactiverefs', () => {
     expectClickLogsLengthToBe(testRefsComponent, 1);
 
     // Begin incrementing clicks (and therefore refs).
-    clickIncrementer.click();
+    await act(() => {
+      clickIncrementer.click();
+    });
     expectClickLogsLengthToBe(testRefsComponent, 2);
 
-    clickIncrementer.click();
+    await act(() => {
+      clickIncrementer.click();
+    });
     expectClickLogsLengthToBe(testRefsComponent, 3);
 
     // Now reset again
-    testRefsComponent.refs.resetDiv.click();
+    await act(() => {
+      testRefsComponent.refs.resetDiv.click();
+    });
     expectClickLogsLengthToBe(testRefsComponent, 1);
   });
 });
@@ -218,9 +235,10 @@ describe('ref swapping', () => {
   beforeEach(() => {
     jest.resetModules();
     React = require('react');
-    ReactDOM = require('react-dom');
+    ReactDOMClient = require('react-dom/client');
     ReactFeatureFlags = require('shared/ReactFeatureFlags');
     ReactTestUtils = require('react-dom/test-utils');
+    act = require('internal-test-utils').act;
 
     RefHopsAround = class extends React.Component {
       state = {count: 0};
@@ -301,18 +319,23 @@ describe('ref swapping', () => {
     expect(refHopsAround.divThreeRef.current).toEqual(thirdDiv);
   });
 
-  it('always has a value for this.refs', () => {
+  it('always has a value for this.refs', async () => {
     class Component extends React.Component {
       render() {
         return <div />;
       }
     }
 
-    const instance = ReactTestUtils.renderIntoDocument(<Component />);
+    const container = document.createElement('div');
+    const root = ReactDOMClient.createRoot(container);
+    let instance;
+    await act(() => {
+      root.render(<Component ref={current => (instance = current)} />);
+    });
     expect(!!instance.refs).toBe(true);
   });
 
-  it('ref called correctly for stateless component', () => {
+  it('ref called correctly for stateless component', async () => {
     let refCalled = 0;
     function Inner(props) {
       return <a ref={props.saveA} />;
@@ -332,7 +355,12 @@ describe('ref swapping', () => {
       }
     }
 
-    ReactTestUtils.renderIntoDocument(<Outer />);
+    const container = document.createElement('div');
+    const root = ReactDOMClient.createRoot(container);
+    await act(() => {
+      root.render(<Outer />);
+    });
+
     expect(refCalled).toBe(1);
   });
 
@@ -355,58 +383,80 @@ describe('ref swapping', () => {
     expect(a.refs[1].nodeName).toBe('DIV');
   });
 
-  it('provides an error for invalid refs', () => {
-    expect(() => {
-      ReactTestUtils.renderIntoDocument(<div ref={10} />);
-    }).toThrow(
+  it('provides an error for invalid refs', async () => {
+    const container = document.createElement('div');
+    const root = ReactDOMClient.createRoot(container);
+    await expect(async () => {
+      await act(() => {
+        root.render(<div ref={10} />);
+      });
+    }).rejects.toThrow(
       'Expected ref to be a function, a string, an object returned by React.createRef(), or null.',
     );
-    expect(() => {
-      ReactTestUtils.renderIntoDocument(<div ref={true} />);
-    }).toThrow(
+    await expect(async () => {
+      await act(() => {
+        root.render(<div ref={true} />);
+      });
+    }).rejects.toThrow(
       'Expected ref to be a function, a string, an object returned by React.createRef(), or null.',
     );
-    expect(() => {
-      ReactTestUtils.renderIntoDocument(<div ref={Symbol('foo')} />);
-    }).toThrow(
+    await expect(async () => {
+      await act(() => {
+        root.render(<div ref={Symbol('foo')} />);
+      });
+    }).rejects.toThrow(
       'Expected ref to be a function, a string, an object returned by React.createRef(), or null.',
     );
-    // This works
-    ReactTestUtils.renderIntoDocument(<div ref={undefined} />);
-    ReactTestUtils.renderIntoDocument({
-      $$typeof: Symbol.for('react.element'),
-      type: 'div',
-      props: {},
-      key: null,
-      ref: null,
+
+    await act(() => {
+      root.render(<div ref={undefined} />);
     });
-    // But this doesn't
-    expect(() => {
-      ReactTestUtils.renderIntoDocument({
+
+    await act(() => {
+      root.render({
         $$typeof: Symbol.for('react.element'),
         type: 'div',
         props: {},
         key: null,
-        ref: undefined,
+        ref: null,
+      });
+    });
+
+    // But this doesn't
+    await expect(async () => {
+      await act(() => {
+        root.render({
+          $$typeof: Symbol.for('react.element'),
+          type: 'div',
+          props: {},
+          key: null,
+          ref: undefined,
+        });
       });
-    }).toThrow(
+    }).rejects.toThrow(
       'Expected ref to be a function, a string, an object returned by React.createRef(), or null.',
     );
   });
 });
 
 describe('root level refs', () => {
-  it('attaches and detaches root refs', () => {
+  it('attaches and detaches root refs', async () => {
     let inst = null;
 
     // host node
     let ref = jest.fn(value => (inst = value));
     const container = document.createElement('div');
-    let result = ReactDOM.render(<div ref={ref} />, container);
+    let root = ReactDOMClient.createRoot(container);
+    await act(() => {
+      root.render(<div ref={ref} />);
+    });
+    let result = container.firstChild;
     expect(ref).toHaveBeenCalledTimes(1);
     expect(ref.mock.calls[0][0]).toBeInstanceOf(HTMLDivElement);
     expect(result).toBe(ref.mock.calls[0][0]);
-    ReactDOM.unmountComponentAtNode(container);
+    await act(() => {
+      root.unmount();
+    });
     expect(ref).toHaveBeenCalledTimes(2);
     expect(ref.mock.calls[1][0]).toBe(null);
 
@@ -422,17 +472,20 @@ describe('root level refs', () => {
 
     inst = null;
     ref = jest.fn(value => (inst = value));
-    result = ReactDOM.render(<Comp ref={ref} />, container);
+    root = ReactDOMClient.createRoot(container);
+    await act(() => {
+      root.render(<Comp ref={ref} />);
+    });
 
     expect(ref).toHaveBeenCalledTimes(1);
     expect(inst).toBeInstanceOf(Comp);
-    expect(result).toBe(inst);
 
     // ensure we have the correct instance
-    expect(result.method()).toBe(true);
     expect(inst.method()).toBe(true);
 
-    ReactDOM.unmountComponentAtNode(container);
+    await act(() => {
+      root.unmount();
+    });
     expect(ref).toHaveBeenCalledTimes(2);
     expect(ref.mock.calls[1][0]).toBe(null);
 
@@ -441,38 +494,45 @@ describe('root level refs', () => {
     ref = jest.fn(value => (inst = value));
     let divInst = null;
     const ref2 = jest.fn(value => (divInst = value));
-    result = ReactDOM.render(
-      [
+    root = ReactDOMClient.createRoot(container);
+    await act(() => {
+      root.render([
         <Comp ref={ref} key="a" />,
         5,
         <div ref={ref2} key="b">
           Hello
         </div>,
-      ],
-      container,
-    );
+      ]);
+    });
 
     // first call should be `Comp`
     expect(ref).toHaveBeenCalledTimes(1);
     expect(ref.mock.calls[0][0]).toBeInstanceOf(Comp);
-    expect(result).toBe(ref.mock.calls[0][0]);
 
     expect(ref2).toHaveBeenCalledTimes(1);
     expect(divInst).toBeInstanceOf(HTMLDivElement);
-    expect(result).not.toBe(divInst);
 
-    ReactDOM.unmountComponentAtNode(container);
+    await act(() => {
+      root.unmount();
+    });
     expect(ref).toHaveBeenCalledTimes(2);
     expect(ref.mock.calls[1][0]).toBe(null);
     expect(ref2).toHaveBeenCalledTimes(2);
     expect(ref2.mock.calls[1][0]).toBe(null);
 
     // null
-    result = ReactDOM.render(null, container);
+    root = ReactDOMClient.createRoot(container);
+    await act(() => {
+      root.render(null);
+    });
+    result = container.firstChild;
     expect(result).toBe(null);
 
     // primitives
-    result = ReactDOM.render(5, container);
+    await act(() => {
+      root.render(5);
+    });
+    result = container.firstChild;
     expect(result).toBeInstanceOf(Text);
   });
 });
@@ -489,12 +549,15 @@ describe('creating element with string ref in constructor', () => {
     }
   }
 
-  it('throws an error', () => {
-    ReactTestUtils = require('react-dom/test-utils');
+  it('throws an error', async () => {
+    await expect(async function () {
+      const container = document.createElement('div');
+      const root = ReactDOMClient.createRoot(container);
 
-    expect(function () {
-      ReactTestUtils.renderIntoDocument(<RefTest />);
-    }).toThrowError(
+      await act(() => {
+        root.render(<RefTest />);
+      });
+    }).rejects.toThrowError(
       'Element ref was specified as a string (p) but no owner was set. This could happen for one of' +
         ' the following reasons:\n' +
         '1. You may be adding a ref to a function component\n' +
@@ -506,7 +569,7 @@ describe('creating element with string ref in constructor', () => {
 });
 
 describe('strings refs across renderers', () => {
-  it('does not break', () => {
+  it('does not break', async () => {
     class Parent extends React.Component {
       render() {
         // This component owns both refs.
@@ -524,7 +587,11 @@ describe('strings refs across renderers', () => {
         // One ref is being rendered later using another renderer copy.
         jest.resetModules();
         const AnotherCopyOfReactDOM = require('react-dom');
-        AnotherCopyOfReactDOM.render(this.props.child2, div2);
+        const AnotherCopyOfReactDOMClient = require('react-dom/client');
+        const root = AnotherCopyOfReactDOMClient.createRoot(div2);
+        AnotherCopyOfReactDOM.flushSync(() => {
+          root.render(this.props.child2);
+        });
       }
       render() {
         // The other one is being rendered directly.
@@ -535,9 +602,20 @@ describe('strings refs across renderers', () => {
     const div1 = document.createElement('div');
     const div2 = document.createElement('div');
 
+    const root = ReactDOMClient.createRoot(div1);
     let inst;
-    expect(() => {
-      inst = ReactDOM.render(<Parent />, div1);
+    await expect(async () => {
+      await act(() => {
+        root.render(
+          <Parent
+            ref={current => {
+              if (current !== null) {
+                inst = current;
+              }
+            }}
+          />,
+        );
+      });
     }).toErrorDev([
       'Warning: Component "Indirection" contains the string ref "child1". ' +
         'Support for string refs will be removed in a future major release. ' +
@@ -551,9 +629,11 @@ describe('strings refs across renderers', () => {
     expect(inst.refs.child1.tagName).toBe('DIV');
     expect(inst.refs.child1).toBe(div1.firstChild);
 
-    expect(() => {
+    await expect(async () => {
       // Now both refs should be rendered.
-      ReactDOM.render(<Parent />, div1);
+      await act(() => {
+        root.render(<Parent />);
+      });
     }).toErrorDev(
       [
         'Warning: Component "Root" contains the string ref "child2". ' +
@@ -571,35 +651,41 @@ describe('strings refs across renderers', () => {
 });
 
 describe('refs return clean up function', () => {
-  it('calls clean up function if it exists', () => {
+  it('calls clean up function if it exists', async () => {
     const container = document.createElement('div');
     let cleanUp = jest.fn();
     let setup = jest.fn();
 
-    ReactDOM.render(
-      <div
-        ref={_ref => {
-          setup(_ref);
-          return cleanUp;
-        }}
-      />,
-      container,
-    );
+    const root = ReactDOMClient.createRoot(container);
 
-    ReactDOM.render(
-      <div
-        ref={_ref => {
-          setup(_ref);
-        }}
-      />,
-      container,
-    );
+    await act(() => {
+      root.render(
+        <div
+          ref={_ref => {
+            setup(_ref);
+            return cleanUp;
+          }}
+        />,
+      );
+    });
+
+    await act(() => {
+      root.render(
+        <div
+          ref={_ref => {
+            setup(_ref);
+          }}
+        />,
+      );
+    });
 
     expect(setup).toHaveBeenCalledTimes(2);
     expect(cleanUp).toHaveBeenCalledTimes(1);
     expect(cleanUp.mock.calls[0][0]).toBe(undefined);
 
-    ReactDOM.render(<div ref={_ref => {}} />, container);
+    await act(() => {
+      root.render(<div ref={_ref => {}} />);
+    });
 
     expect(cleanUp).toHaveBeenCalledTimes(1);
     expect(setup).toHaveBeenCalledTimes(3);
@@ -608,34 +694,36 @@ describe('refs return clean up function', () => {
     cleanUp = jest.fn();
     setup = jest.fn();
 
-    ReactDOM.render(
-      <div
-        ref={_ref => {
-          setup(_ref);
-          return cleanUp;
-        }}
-      />,
-      container,
-    );
+    await act(() => {
+      root.render(
+        <div
+          ref={_ref => {
+            setup(_ref);
+            return cleanUp;
+          }}
+        />,
+      );
+    });
 
     expect(setup).toHaveBeenCalledTimes(1);
     expect(cleanUp).toHaveBeenCalledTimes(0);
 
-    ReactDOM.render(
-      <div
-        ref={_ref => {
-          setup(_ref);
-          return cleanUp;
-        }}
-      />,
-      container,
-    );
+    await act(() => {
+      root.render(
+        <div
+          ref={_ref => {
+            setup(_ref);
+            return cleanUp;
+          }}
+        />,
+      );
+    });
 
     expect(setup).toHaveBeenCalledTimes(2);
     expect(cleanUp).toHaveBeenCalledTimes(1);
   });
 
-  it('handles ref functions with stable identity', () => {
+  it('handles ref functions with stable identity', async () => {
     const container = document.createElement('div');
     const cleanUp = jest.fn();
     const setup = jest.fn();
@@ -645,50 +733,38 @@ describe('refs return clean up function', () => {
       return cleanUp;
     }
 
-    ReactDOM.render(<div ref={_onRefChange} />, container);
+    const root = ReactDOMClient.createRoot(container);
+    await act(() => {
+      root.render(<div ref={_onRefChange} />);
+    });
 
     expect(setup).toHaveBeenCalledTimes(1);
     expect(cleanUp).toHaveBeenCalledTimes(0);
 
-    ReactDOM.render(
-      <div className="niceClassName" ref={_onRefChange} />,
-      container,
-    );
+    await act(() => {
+      root.render(<div className="niceClassName" ref={_onRefChange} />);
+    });
 
     expect(setup).toHaveBeenCalledTimes(1);
     expect(cleanUp).toHaveBeenCalledTimes(0);
 
-    ReactDOM.render(<div />, container);
+    await act(() => {
+      root.render(<div />);
+    });
 
     expect(setup).toHaveBeenCalledTimes(1);
     expect(cleanUp).toHaveBeenCalledTimes(1);
   });
 
-  it('warns if clean up function is returned when called with null', () => {
+  it('warns if clean up function is returned when called with null', async () => {
     const container = document.createElement('div');
     const cleanUp = jest.fn();
     const setup = jest.fn();
     let returnCleanUp = false;
 
-    ReactDOM.render(
-      <div
-        ref={_ref => {
-          setup(_ref);
-          if (returnCleanUp) {
-            return cleanUp;
-          }
-        }}
-      />,
-      container,
-    );
-
-    expect(setup).toHaveBeenCalledTimes(1);
-    expect(cleanUp).toHaveBeenCalledTimes(0);
-
-    returnCleanUp = true;
-
-    expect(() => {
-      ReactDOM.render(
+    const root = ReactDOMClient.createRoot(container);
+    await act(() => {
+      root.render(
         <div
           ref={_ref => {
             setup(_ref);
@@ -697,8 +773,27 @@ describe('refs return clean up function', () => {
             }
           }}
         />,
-        container,
       );
+    });
+
+    expect(setup).toHaveBeenCalledTimes(1);
+    expect(cleanUp).toHaveBeenCalledTimes(0);
+
+    returnCleanUp = true;
+
+    await expect(async () => {
+      await act(() => {
+        root.render(
+          <div
+            ref={_ref => {
+              setup(_ref);
+              if (returnCleanUp) {
+                return cleanUp;
+              }
+            }}
+          />,
+        );
+      });
     }).toErrorDev('Unexpected return value from a callback ref in div');
   });
 });
diff --git a/packages/react-dom/src/__tests__/refsLegacy-test.js b/packages/react-dom/src/__tests__/refsLegacy-test.js
new file mode 100644
index 0000000000000..c3af817892ffa
--- /dev/null
+++ b/packages/react-dom/src/__tests__/refsLegacy-test.js
@@ -0,0 +1,101 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @emails react-core
+ */
+
+'use strict';
+
+let React = require('react');
+let ReactDOM = require('react-dom');
+
+describe('root level refs with legacy APIs', () => {
+  beforeEach(() => {
+    jest.resetModules();
+    React = require('react');
+    ReactDOM = require('react-dom');
+  });
+
+  it('attaches and detaches root refs', () => {
+    let inst = null;
+
+    // host node
+    let ref = jest.fn(value => (inst = value));
+    const container = document.createElement('div');
+    let result = ReactDOM.render(<div ref={ref} />, container);
+    expect(ref).toHaveBeenCalledTimes(1);
+    expect(ref.mock.calls[0][0]).toBeInstanceOf(HTMLDivElement);
+    expect(result).toBe(ref.mock.calls[0][0]);
+    ReactDOM.unmountComponentAtNode(container);
+    expect(ref).toHaveBeenCalledTimes(2);
+    expect(ref.mock.calls[1][0]).toBe(null);
+
+    // composite
+    class Comp extends React.Component {
+      method() {
+        return true;
+      }
+      render() {
+        return <div>Comp</div>;
+      }
+    }
+
+    inst = null;
+    ref = jest.fn(value => (inst = value));
+    result = ReactDOM.render(<Comp ref={ref} />, container);
+
+    expect(ref).toHaveBeenCalledTimes(1);
+    expect(inst).toBeInstanceOf(Comp);
+    expect(result).toBe(inst);
+
+    // ensure we have the correct instance
+    expect(result.method()).toBe(true);
+    expect(inst.method()).toBe(true);
+
+    ReactDOM.unmountComponentAtNode(container);
+    expect(ref).toHaveBeenCalledTimes(2);
+    expect(ref.mock.calls[1][0]).toBe(null);
+
+    // fragment
+    inst = null;
+    ref = jest.fn(value => (inst = value));
+    let divInst = null;
+    const ref2 = jest.fn(value => (divInst = value));
+    result = ReactDOM.render(
+      [
+        <Comp ref={ref} key="a" />,
+        5,
+        <div ref={ref2} key="b">
+          Hello
+        </div>,
+      ],
+      container,
+    );
+
+    // first call should be `Comp`
+    expect(ref).toHaveBeenCalledTimes(1);
+    expect(ref.mock.calls[0][0]).toBeInstanceOf(Comp);
+    expect(result).toBe(ref.mock.calls[0][0]);
+
+    expect(ref2).toHaveBeenCalledTimes(1);
+    expect(divInst).toBeInstanceOf(HTMLDivElement);
+    expect(result).not.toBe(divInst);
+
+    ReactDOM.unmountComponentAtNode(container);
+    expect(ref).toHaveBeenCalledTimes(2);
+    expect(ref.mock.calls[1][0]).toBe(null);
+    expect(ref2).toHaveBeenCalledTimes(2);
+    expect(ref2.mock.calls[1][0]).toBe(null);
+
+    // null
+    result = ReactDOM.render(null, container);
+    expect(result).toBe(null);
+
+    // primitives
+    result = ReactDOM.render(5, container);
+    expect(result).toBeInstanceOf(Text);
+  });
+});