[core] Prepare focus visible polyfill in ref phase

packages/material-ui/src/ButtonBase/ButtonBase.js

@@ -110,20 +110,7 @@ const ButtonBase = React.forwardRef(function ButtonBase(props, ref) {
   if (disabled && focusVisible) {
     setFocusVisible(false);
   }
-  const getOwnerDocument = React.useCallback(() => {
-    const button = getButtonNode();
-    if (button == null) {
-      throw new Error(
-        [
-          `Material-UI: expected an Element but found ${button}.`,
-          'Please check your console for additional warnings and try fixing those.',
-          'If the error persists please file an issue.',
-        ].join(' '),
-      );
-    }
-    return button.ownerDocument;
-  }, []);
-  const { isFocusVisible, onBlurVisible } = useIsFocusVisible(getOwnerDocument);
+  const { isFocusVisible, onBlurVisible, ref: focusVisibleRef } = useIsFocusVisible();
 
   React.useImperativeHandle(
     action,
@@ -281,7 +268,8 @@ const ButtonBase = React.forwardRef(function ButtonBase(props, ref) {
   }
 
   const handleUserRef = useForkRef(buttonRefProp, ref);
-  const handleRef = useForkRef(handleUserRef, buttonRef);
+  const handleOwnRef = useForkRef(focusVisibleRef, buttonRef);
+  const handleRef = useForkRef(handleUserRef, handleOwnRef);
 
   return (
     <ComponentProp

packages/material-ui/src/ButtonBase/ButtonBase.test.js

@@ -637,7 +637,7 @@ describe('<ButtonBase />', () => {
       PropTypes.resetWarningCache();
     });
 
-    it('throws with additional warnings on invalid `component` prop', () => {
+    it('warns on invalid `component` prop', () => {
       // Only run the test on node. On the browser the thrown error is not caught
       if (!/jsdom/.test(window.navigator.userAgent)) {
         return;
@@ -648,19 +648,12 @@ describe('<ButtonBase />', () => {
       }
 
       // cant match the error message here because flakiness with mocha watchmode
-      assert.throws(() => mount(<ButtonBase component={Component} />));
+      mount(<ButtonBase component={Component} />);
 
       assert.include(
         consoleErrorMock.args()[0][0],
         'Invalid prop `component` supplied to `ForwardRef(ButtonBase)`. Expected an element type that can hold a ref',
       );
-      // first mount includes React warning that isn't logged on subsequent calls
-      // in watchmode because it's cached
-      const customErrorIndex = consoleErrorMock.callCount() === 3 ? 1 : 2;
-      assert.include(
-        consoleErrorMock.args()[customErrorIndex][0],
-        'Error: Material-UI: expected an Element but found null. Please check your console for additional warnings and try fixing those.',
-      );
     });
   });
 });

packages/material-ui/src/Tooltip/Tooltip.js

@@ -113,12 +113,6 @@ function Tooltip(props) {
   const enterTimer = React.useRef();
   const leaveTimer = React.useRef();
   const touchTimer = React.useRef();
-  // can be removed once we drop support for non ref forwarding class components
-  const handleOwnRef = React.useCallback(instance => {
-    // #StrictMode ready
-    setChildNode(ReactDOM.findDOMNode(instance));
-  }, []);
-  const handleRef = useForkRef(children.ref, handleOwnRef);
 
   React.useEffect(() => {
     warning(
@@ -205,13 +199,7 @@ function Tooltip(props) {
     }
   };
 
-  const getOwnerDocument = React.useCallback(() => {
-    if (childNode == null) {
-      return null;
-    }
-    return childNode.ownerDocument;
-  }, [childNode]);
-  const { isFocusVisible, onBlurVisible } = useIsFocusVisible(getOwnerDocument);
+  const { isFocusVisible, onBlurVisible, ref: focusVisibleRef } = useIsFocusVisible();
   const [childIsFocusVisible, setChildIsFocusVisible] = React.useState(false);
   function handleBlur() {
     if (childIsFocusVisible) {
@@ -310,6 +298,16 @@ function Tooltip(props) {
     }, leaveTouchDelay);
   };
 
+  // can be removed once we drop support for non ref forwarding class components
+  const handleOwnRef = useForkRef(
+    React.useCallback(instance => {
+      // #StrictMode ready
+      setChildNode(ReactDOM.findDOMNode(instance));
+    }, []),
+    focusVisibleRef,
+  );
+  const handleRef = useForkRef(children.ref, handleOwnRef);
+
   let open = isControlled ? openProp : openState;
 
   // There is no point in displaying an empty tooltip.

packages/material-ui/src/utils/focusVisible.js

@@ -1,5 +1,6 @@
 // based on https://github.com/WICG/focus-visible/blob/v4.1.5/src/focus-visible.js
 import React from 'react';
+import ReactDOM from 'react-dom';
 
 let hadKeyboardEvent = true;
 let hadFocusVisibleRecently = false;
@@ -122,13 +123,13 @@ export function handleBlurVisible() {
   }, 100);
 }
 
-export function useIsFocusVisible(getOwnerDocument) {
-  React.useEffect(() => {
-    const ownerDocument = getOwnerDocument();
-    if (ownerDocument != null) {
-      prepare(ownerDocument);
+export function useIsFocusVisible() {
+  const ref = React.useCallback(instance => {
+    const node = ReactDOM.findDOMNode(instance);
+    if (node != null) {
+      prepare(node.ownerDocument);
     }
-  }, [getOwnerDocument]);
+  }, []);
 
-  return { isFocusVisible, onBlurVisible: handleBlurVisible };
+  return { isFocusVisible, onBlurVisible: handleBlurVisible, ref };
 }

packages/material-ui/src/utils/focusVisible.test.js

@@ -17,17 +17,9 @@ function simulatePointerDevice() {
 }
 
 const SimpleButton = React.forwardRef(function SimpleButton(props, ref) {
-  const [element, setElement] = React.useState(null);
-  const getOwnerDocument = React.useCallback(() => {
-    if (element === null) {
-      return null;
-    }
-
-    return element.ownerDocument;
-  }, [element]);
-  const { isFocusVisible, onBlurVisible } = useIsFocusVisible(getOwnerDocument);
+  const { isFocusVisible, onBlurVisible, ref: focusVisibleRef } = useIsFocusVisible();
 
-  const handleRef = useForkRef(setElement, ref);
+  const handleRef = useForkRef(focusVisibleRef, ref);
 
   const [focusVisible, setFocusVisible] = React.useState(false);