From 8c6507fb89887b4802056fa8b50baba7f9d3ae5b Mon Sep 17 00:00:00 2001
From: Morgan Zolob <mzolob@gmail.com>
Date: Mon, 9 Dec 2019 22:04:04 -0800
Subject: [PATCH] Fix passing ref to functional components (#153)

---
 src/index.tsx        | 20 +++++++++----
 tests/basic.test.jsx | 67 ++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 76 insertions(+), 11 deletions(-)

diff --git a/src/index.tsx b/src/index.tsx
index d100aa43..21c6c4fb 100644
--- a/src/index.tsx
+++ b/src/index.tsx
@@ -1,4 +1,4 @@
-import React, { HTMLAttributes } from 'react';
+import React, { HTMLAttributes, RefAttributes, ReactInstance } from 'react';
 import ReactDOM from 'react-dom';
 import contains from 'rc-util/lib/Dom/contains';
 import findDOMNode from 'rc-util/lib/Dom/findDOMNode';
@@ -676,7 +676,8 @@ export function generateTrigger(PortalComponent: any): React.ComponentClass<Trig
       const { popupVisible } = this.state;
       const { children, forceRender, alignPoint, className } = this.props;
       const child = React.Children.only(children) as React.ReactElement;
-      const newChildProps: HTMLAttributes<HTMLElement> & { key: string } = { key: 'trigger' };
+      const newChildProps: HTMLAttributes<HTMLElement> &
+        RefAttributes<ReactInstance> & { key: string } = { key: 'trigger' };
 
       if (this.isContextMenuToShow()) {
         newChildProps.onContextMenu = this.onContextMenu;
@@ -721,10 +722,17 @@ export function generateTrigger(PortalComponent: any): React.ComponentClass<Trig
       if (childrenClassName) {
         newChildProps.className = childrenClassName;
       }
-      const trigger = React.cloneElement(child, {
-        ...newChildProps,
-        ref: composeRef(this.triggerRef, (child as any).ref),
-      });
+
+      // Prevent adding ref to Functional child components
+      if (
+        !child.type ||
+        typeof child.type !== 'function' ||
+        (child.type.prototype && child.type.prototype.isReactComponent)
+      ) {
+        newChildProps.ref = composeRef(this.triggerRef, (child as any).ref);
+      }
+
+      const trigger = React.cloneElement(child, newChildProps);
 
       let portal: React.ReactElement;
       // prevent unmounting after it's rendered
diff --git a/tests/basic.test.jsx b/tests/basic.test.jsx
index befc76c5..2661730f 100644
--- a/tests/basic.test.jsx
+++ b/tests/basic.test.jsx
@@ -514,14 +514,11 @@ describe('Trigger.Basic', () => {
   });
 
   it('support function component', () => {
-    const NoRef = React.forwardRef((props, ref) => {
-      React.useImperativeHandle(ref, () => null);
-      return (
+    const FuncComp = props => (
         <div className="target" {...props}>
           click
         </div>
       );
-    });
 
     const wrapper = mount(
       <Trigger
@@ -529,7 +526,7 @@ describe('Trigger.Basic', () => {
         popupAlign={placementAlignMap.left}
         popup={<strong className="x-content">tooltip2</strong>}
       >
-        <NoRef />
+        <FuncComp />
       </Trigger>,
     );
 
@@ -539,4 +536,64 @@ describe('Trigger.Basic', () => {
     wrapper.trigger();
     expect(wrapper.isHidden()).toBeTruthy();
   });
+
+  describe('passes ref to children where applicable', () => {
+    function getRefResult(component) {
+      const triggerMock = jest.fn();
+      const wrapper = mount(
+        <Trigger
+          action={['click']}
+          popupAlign={placementAlignMap.left}
+          popup={<strong className="x-content">tooltip2</strong>}
+          ref={triggerMock}
+        >
+          {component}
+        </Trigger>,
+      );
+
+      wrapper.trigger();
+
+      return triggerMock.mock.calls[0][0].triggerRef.current;
+    }
+
+    it('does not pass ref to function component', () => {
+      const NoRef = props => (
+        <div id="target" {...props}>
+          click
+        </div>
+      );
+
+      const refVal = getRefResult(<NoRef />);
+      expect(refVal).toBeNull();
+    });
+
+    it('does pass ref to forwardRef function component', () => {
+      const WithRef = React.forwardRef((props, ref) => (
+        <div id="target" {...props} ref={ref}>
+          click
+        </div>
+      ));
+
+      const refVal = getRefResult(<WithRef />);
+      expect(refVal.id).toBe('target');
+    });
+
+    it('does pass ref to class component', () => {
+      class ClassComp extends React.Component {
+        id = 'target';
+
+        render() {
+          return <div>click</div>;
+        }
+      }
+
+      const refVal = getRefResult(<ClassComp />);
+      expect(refVal.id).toBe('target');
+    });
+
+    it('does pass ref to element', () => {
+      const refVal = getRefResult(<div id="target">click</div>);
+      expect(refVal.id).toBe('target');
+    });
+  });
 });