fix(Ref): handle primitive elements with findDOMNode (#35779)

Author: layershifter

packages/fluentui/react-component-ref/src/Ref.tsx

@@ -16,7 +16,7 @@ export interface RefProps {
 }
 
 export interface RefState {
-  kind: 'self' | 'forward' | 'find' | null;
+  kind: 'self' | 'direct' | 'find' | null;
 }
 
 // ========================================================
@@ -51,7 +51,7 @@ function isFiberRef(node: Element | Fiber | Text | null): boolean {
 
 export class Ref extends React.Component<RefProps, RefState> {
   nodeForFind?: Node | null;
-  nodeForForward?: Node | null;
+  nodeForDirect?: Node | null;
 
   state = { kind: null };
 
@@ -62,20 +62,20 @@ export class Ref extends React.Component<RefProps, RefState> {
       return { kind: 'self' };
     }
 
-    if (ReactIs.isForwardRef(child)) {
-      return { kind: 'forward' };
+    if (typeof child.type === 'string' || ReactIs.isForwardRef(child)) {
+      return { kind: 'direct' };
     }
 
     return { kind: 'find' };
   }
 
-  handleForwardRefOverride = (node: HTMLElement) => {
+  handleRefOverride = (node: HTMLElement) => {
     const { children, innerRef } = this.props;
 
     handleRef((children as React.ReactElement & { ref: React.Ref<any> }).ref, node);
     handleRef(innerRef, node);
 
-    this.nodeForForward = node;
+    this.nodeForDirect = node;
   };
 
   handleSelfOverride = (node: HTMLElement) => {
@@ -101,9 +101,9 @@ export class Ref extends React.Component<RefProps, RefState> {
   }
 
   componentDidUpdate(prevProps: RefProps) {
-    if (this.state.kind === 'forward') {
+    if (this.state.kind === 'direct') {
       if (prevProps.innerRef !== this.props.innerRef) {
-        handleRef(this.props.innerRef, this.nodeForForward);
+        handleRef(this.props.innerRef, this.nodeForDirect);
       }
     } else if (this.state.kind === 'find') {
       let currentNode = ReactDOM.findDOMNode(this);
@@ -142,9 +142,9 @@ export class Ref extends React.Component<RefProps, RefState> {
       return childWithProps;
     }
 
-    if (this.state.kind === 'forward') {
+    if (this.state.kind === 'direct') {
       return React.cloneElement(childWithProps, {
-        ref: this.handleForwardRefOverride,
+        ref: this.handleRefOverride,
       });
     }
 

packages/fluentui/react-component-ref/test/Ref-test.tsx

@@ -61,15 +61,18 @@ describe('Ref', () => {
     });
 
     it('passes an updated node', () => {
+      const ComponentA = () => <div />;
+      const ComponentB = () => <button />;
+
       const innerRef = jest.fn();
       const wrapper = mount(
         <Ref innerRef={innerRef}>
-          <div />
+          <ComponentA />
         </Ref>,
       );
 
       expect(innerRef).toHaveBeenCalledWith(expect.objectContaining({ tagName: 'DIV' }));
-      wrapper.setProps({ children: <button /> });
+      wrapper.setProps({ children: <ComponentB /> });
 
       expect(innerRef).toHaveBeenCalledTimes(2);
       expect(innerRef).toHaveBeenCalledWith(expect.objectContaining({ tagName: 'BUTTON' }));
@@ -110,36 +113,62 @@ describe('Ref', () => {
     });
 
     it('resolves once on node/ref change', () => {
+      const ComponentA = () => <div />;
+      const ComponentB = () => <span />;
+
       const initialRef = jest.fn();
       const updatedRef = jest.fn();
+
       const wrapper = mount(
         <Ref innerRef={initialRef}>
-          <div />
+          <ComponentA />
         </Ref>,
       );
 
       expect(initialRef).toHaveBeenCalledTimes(1);
       expect(updatedRef).not.toHaveBeenCalled();
 
+      // ---
+
       jest.resetAllMocks();
-      wrapper.setProps({ children: <span />, innerRef: updatedRef });
+
+      wrapper.setProps({ children: <ComponentB />, innerRef: updatedRef });
 
       expect(initialRef).not.toHaveBeenCalled();
       expect(updatedRef).toHaveBeenCalledTimes(1);
     });
 
     it('always returns "null" for react-test-renderer', () => {
+      const Component = () => <div />;
       const innerRef = jest.fn();
-      const ref = jest.fn();
 
       create(
         <Ref innerRef={innerRef}>
-          <div ref={ref} />
+          <Component />
         </Ref>,
       );
 
       expect(innerRef).toHaveBeenCalledWith(null);
-      expect(ref).toHaveBeenCalledWith(null);
+    });
+
+    it('handles unmount', () => {
+      const Component = () => <div />;
+      const innerRef = jest.fn();
+
+      const wrapper = mount(
+        <Ref innerRef={innerRef}>
+          <Component />
+        </Ref>,
+      );
+
+      expect(innerRef).toHaveBeenCalledWith(expect.objectContaining({ tagName: 'DIV' }));
+
+      // ---
+
+      jest.resetAllMocks();
+      wrapper.unmount();
+
+      expect(innerRef).toHaveBeenCalledWith(null);
     });
   });
 
@@ -201,6 +230,65 @@ describe('Ref', () => {
     });
   });
 
+  describe('kind="forward" (plain element)', () => {
+    it('applies refs', () => {
+      const innerRef = jest.fn();
+      const elRef = jest.fn();
+
+      mount(
+        <Ref innerRef={innerRef}>
+          <button ref={elRef} />
+        </Ref>,
+      );
+
+      expect(innerRef).toHaveBeenCalledWith(expect.objectContaining({ tagName: 'BUTTON' }));
+      expect(elRef).toHaveBeenCalledWith(expect.objectContaining({ tagName: 'BUTTON' }));
+    });
+
+    it('passes an updated node', () => {
+      const innerRef = jest.fn();
+
+      mount(
+        <Ref innerRef={innerRef}>
+          <button />
+        </Ref>,
+      );
+
+      expect(innerRef).toHaveBeenCalledWith(expect.objectContaining({ tagName: 'BUTTON' }));
+
+      // ---
+
+      jest.resetAllMocks();
+
+      mount(
+        <Ref innerRef={innerRef}>
+          <div />
+        </Ref>,
+      );
+
+      expect(innerRef).toHaveBeenCalledWith(expect.objectContaining({ tagName: 'DIV' }));
+    });
+
+    it('handles unmount', () => {
+      const innerRef = jest.fn();
+
+      const wrapper = mount(
+        <Ref innerRef={innerRef}>
+          <button />
+        </Ref>,
+      );
+
+      expect(innerRef).toHaveBeenCalledWith(expect.objectContaining({ tagName: 'BUTTON' }));
+
+      // ---
+
+      jest.resetAllMocks();
+      wrapper.unmount();
+
+      expect(innerRef).toHaveBeenCalledWith(null);
+    });
+  });
+
   describe('kind="self"', () => {
     it('works with "forwardRef" API', () => {
       const forwardedRef = React.createRef<HTMLButtonElement>();
@@ -239,7 +327,10 @@ describe('Ref', () => {
       expect(innerRef).toHaveBeenCalledWith(expect.objectContaining({ tagName: 'DIV' }));
       expect(outerRef).toHaveBeenCalledWith(expect.objectContaining({ tagName: 'DIV' }));
 
+      // ---
+
       jest.resetAllMocks();
+
       wrapper.setProps({
         children: (
           <Ref innerRef={innerRef}>