[RootRef] Apply the same logic than with React Ref

mui · Jul 28, 2018 · 7dfc7a6 · 7dfc7a6
1 parent ee9bc65
commit 7dfc7a6
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 21 deletions.
diff --git a/packages/material-ui/src/RootRef/RootRef.js b/packages/material-ui/src/RootRef/RootRef.js
@@ -3,6 +3,14 @@ import ReactDOM from 'react-dom';
 import PropTypes from 'prop-types';
 import exactProp from '../utils/exactProp';
 
+function setRef(ref, value) {
+  if (typeof ref === 'function') {
+    ref(value);
+  } else if (ref) {
+    ref.current = value;
+  }
+}
+
 /**
  * Helper component to allow attaching a ref to a
  * wrapped element to access the underlying DOM element.
@@ -35,22 +43,18 @@ import exactProp from '../utils/exactProp';
  */
 class RootRef extends React.Component {
   componentDidMount() {
-    const rootRef = this.props.rootRef;
-    const node = ReactDOM.findDOMNode(this);
-    if (typeof rootRef === 'function') {
-      rootRef(node);
-    } else if (rootRef) {
-      rootRef.current = node;
+    setRef(this.props.rootRef, ReactDOM.findDOMNode(this));
+  }
+
+  componentDidUpdate(prevProps) {
+    if (prevProps.rootRef !== this.props.rootRef) {
+      setRef(prevProps.rootRef, null);
+      setRef(this.props.rootRef, ReactDOM.findDOMNode(this));
     }
   }
 
   componentWillUnmount() {
-    const rootRef = this.props.rootRef;
-    if (typeof rootRef === 'function') {
-      rootRef(null);
-    } else if (rootRef) {
-      rootRef.current = null;
-    }
+    setRef(this.props.rootRef, null);
   }
 
   render() {

diff --git a/packages/material-ui/src/RootRef/RootRef.test.js b/packages/material-ui/src/RootRef/RootRef.test.js
@@ -3,6 +3,8 @@ import { assert } from 'chai';
 import { createMount } from '../test-utils';
 import RootRef from './RootRef';
 
+const Fn = () => <div />;
+
 describe('<RootRef />', () => {
   let mount;
 
@@ -15,7 +17,6 @@ describe('<RootRef />', () => {
   });
 
   it('call rootRef function on mount and unmount', () => {
-    const Fn = () => <div />;
     const results = [];
     const wrapper = mount(
       <RootRef rootRef={ref => results.push(ref)}>
@@ -30,7 +31,6 @@ describe('<RootRef />', () => {
   });
 
   it('set rootRef current field on mount and unmount', () => {
-    const Fn = () => <div />;
     const ref = React.createRef();
     const wrapper = mount(
       <RootRef rootRef={ref}>
@@ -41,4 +41,24 @@ describe('<RootRef />', () => {
     wrapper.unmount();
     assert.strictEqual(ref.current, null);
   });
+
+  it('should support providing a new rootRef', () => {
+    const results = [];
+    const wrapper = mount(
+      <RootRef rootRef={ref => results.push(ref)}>
+        <Fn />
+      </RootRef>,
+    );
+    assert.strictEqual(results.length, 1);
+    assert.strictEqual(results[0] instanceof window.HTMLDivElement, true);
+    wrapper.setProps({
+      rootRef: ref => results.push(ref),
+    });
+    assert.strictEqual(results.length, 3);
+    assert.strictEqual(results[1], null);
+    assert.strictEqual(results[2] instanceof window.HTMLDivElement, true);
+    wrapper.unmount();
+    assert.strictEqual(results.length, 4);
+    assert.strictEqual(results[3], null);
+  });
 });
diff --git a/packages/material-ui/src/Tooltip/Tooltip.js b/packages/material-ui/src/Tooltip/Tooltip.js
@@ -131,6 +131,10 @@ class Tooltip extends React.Component {
     clearTimeout(this.closeTimer);
   }
 
+  onRootRef = ref => {
+    this.childrenRef = ref;
+  };
+
   handleEnter = event => {
     const { children, enterDelay } = this.props;
     const childrenProps = children.props;
@@ -319,13 +323,7 @@ class Tooltip extends React.Component {
 
     return (
       <React.Fragment>
-        <RootRef
-          rootRef={ref => {
-            this.childrenRef = ref;
-          }}
-        >
-          {React.cloneElement(children, childrenProps)}
-        </RootRef>
+        <RootRef rootRef={this.onRootRef}>{React.cloneElement(children, childrenProps)}</RootRef>
         <Popper
           className={classes.popper}
           placement={placement}