[RootRef] Apply the same logic than with React Ref (#12311)

.size-limit.js

@@ -21,13 +21,13 @@ module.exports = [
  name: 'The initial cost people pay for using one component',
  webpack: true,
  path: 'packages/material-ui/build/Paper/index.js',
- limit: '17.9 KB',
+ limit: '17.7 KB',
  },
  {
  name: 'The size of all the modules of material-ui.',
  webpack: true,
  path: 'packages/material-ui/build/index.js',
- limit: '95.1 KB',
+ limit: '95.2 KB',
  },
  {
  name: 'The main bundle of the docs',

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() {

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);
+ });
 });

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}