fix(ui-motion,ui-portal): fix ref for Transition and Portal

packages/ui-motion/src/Transition/BaseTransition/index.ts

@@ -22,9 +22,9 @@
  * SOFTWARE.
  */
 
-import React, { createRef, ReactElement } from 'react'
+import React, { ReactElement } from 'react'
 
-import { getClassList } from '@instructure/ui-dom-utils'
+import { getClassList, findDOMNode } from '@instructure/ui-dom-utils'
 import {
   ensureSingleChild,
   safeCloneElement
@@ -77,6 +77,18 @@ class BaseTransition extends React.Component<
     transitioning: false
   }
 
+  ref: Element | null = null
+
+  handleRef = (el: Element | null) => {
+    const { elementRef } = this.props
+
+    this.ref = el
+
+    if (typeof elementRef === 'function') {
+      elementRef(el)
+    }
+  }
+
   componentDidMount() {
     this.startTransition(this.props.in, this.props.transitionOnMount)
   }
@@ -139,7 +151,7 @@ class BaseTransition extends React.Component<
 
     const { onTransition } = this.props
 
-    const classList = getClassList(this)
+    const classList = getClassList(this.ref)
 
     const transitionClassName = this.getTransitionClassName(toState)
     const prevTransitionClassName = this.getTransitionClassName(fromState)
@@ -180,7 +192,7 @@ class BaseTransition extends React.Component<
     this.setState({ transitioning: false }, () => {
       if (this._unmounted) return
 
-      const classList = getClassList(this)
+      const classList = getClassList(this.ref)
 
       ;(Object.values(STATES) as BaseTransitionStatesType[]).forEach(
         (state) => {
@@ -314,15 +326,17 @@ class BaseTransition extends React.Component<
     }
   }
 
-  ref = createRef()
-
   renderChildren() {
     return this.props.children
       ? safeCloneElement(
           ensureSingleChild(this.props.children) as ReactElement,
           {
             'aria-hidden': !this.props.in ? true : undefined,
-            ref: this.ref
+            ref: (el: React.ReactInstance | null) => {
+              const ref = (findDOMNode(el) as Element) || null
+
+              this.handleRef(ref)
+            }
           }
         )
       : null

packages/ui-motion/src/Transition/BaseTransition/props.ts

@@ -98,6 +98,11 @@ type TransitionCommonProps = {
    * A single element to animate in and out
    */
   children?: React.ReactNode
+
+  /**
+   * provides a reference to the underlying html root element
+   */
+  elementRef?: (element: Element | null) => void
 }
 
 type OwnProps = {
@@ -171,7 +176,8 @@ const transitionCommonPropTypes: PropValidators<CommonPropKeys> = {
   onExit: PropTypes.func,
   onExiting: PropTypes.func,
   onExited: PropTypes.func,
-  children: PropTypes.node
+  children: PropTypes.node,
+  elementRef: PropTypes.func
 }
 
 const propTypes: PropValidators<PropKeys> = {
@@ -208,7 +214,8 @@ const allowedProps: AllowedPropKeys = [
   'onExiting',
   'onExited',
   'children',
-  'className'
+  'className',
+  'elementRef'
 ]
 
 type BaseTransitionState = {

packages/ui-motion/src/Transition/__tests__/Transition.test.tsx

@@ -22,7 +22,7 @@
  * SOFTWARE.
  */
 
-import React from 'react'
+import React, { Component } from 'react'
 import {
   expect,
   mount,
@@ -45,6 +45,12 @@ const getClass = (
   return classNames[phase]
 }
 
+class ExampleComponent extends Component {
+  render() {
+    return <div>Example</div>
+  }
+}
+
 describe('<Transition />', async () => {
   const types: TransitionType[] = [
     'fade',
@@ -56,7 +62,7 @@ describe('<Transition />', async () => {
   ]
 
   const expectTypeClass = function (type: TransitionType) {
-    it(`should correctly apply classes for '${type}'`, async () => {
+    it(`should correctly apply classes for '${type}' with html element`, async () => {
       const subject = await mount(
         <Transition type={type} in={true}>
           <div>hello</div>
@@ -67,6 +73,18 @@ describe('<Transition />', async () => {
 
       expect(transition.hasClass(getClass(type, 'entered'))).to.be.true()
     })
+
+    it(`should correctly apply classes for '${type}' with Component`, async () => {
+      const subject = await mount(
+        <Transition type={type} in={true}>
+          <ExampleComponent />
+        </Transition>
+      )
+
+      const transition = within(subject.getDOMNode())
+
+      expect(transition.hasClass(getClass(type, 'entered'))).to.be.true()
+    })
   }
 
   types.forEach((type) => {

packages/ui-motion/src/Transition/index.tsx

@@ -67,6 +67,18 @@ class Transition extends Component<TransitionProps> {
 
   static states = BaseTransition.states
 
+  ref: Element | null = null
+
+  handleRef = (el: Element | null) => {
+    const { elementRef } = this.props
+
+    this.ref = el
+
+    if (typeof elementRef === 'function') {
+      elementRef(el)
+    }
+  }
+
   componentDidMount() {
     this.props.makeStyles?.()
   }
@@ -117,6 +129,7 @@ class Transition extends Component<TransitionProps> {
           enteringClassName={styles!.classNames.entering}
           onEntered={this.handleEntered}
           onExited={this.handleExited}
+          elementRef={this.handleRef}
         >
           {children}
         </BaseTransition>

packages/ui-motion/src/Transition/props.ts

@@ -88,7 +88,8 @@ const allowedProps: AllowedPropKeys = [
   'onEntered',
   'onExit',
   'onExiting',
-  'onExited'
+  'onExited',
+  'elementRef'
 ]
 
 export type { TransitionProps, TransitionType, TransitionStyle }

packages/ui-portal/src/Portal/index.tsx

@@ -68,6 +68,19 @@ class Portal extends Component<PortalProps, PortalState> {
 
   DOMNode: PortalNode = null
 
+  ref: PortalNode = null
+
+  handleRef = (el: PortalNode) => {
+    const { elementRef } = this.props
+
+    this.ref = el
+    this.DOMNode = el
+
+    if (typeof elementRef === 'function') {
+      elementRef(el)
+    }
+  }
+
   componentDidMount() {
     if (!canUseDOM) {
       return
@@ -127,11 +140,8 @@ class Portal extends Component<PortalProps, PortalState> {
       typeof this.DOMNode.parentNode.removeChild === 'function'
     ) {
       this.DOMNode.parentNode.removeChild(this.DOMNode)
-      this.DOMNode = null
 
-      if (typeof this.props.elementRef === 'function') {
-        this.props.elementRef(this.DOMNode)
-      }
+      this.handleRef(null)
     }
   }
 
@@ -160,15 +170,11 @@ class Portal extends Component<PortalProps, PortalState> {
         node.setAttribute(name, attributes[name] as string)
       })
 
-      if (typeof elementRef === 'function') {
-        elementRef(node)
-      }
-
-      this.DOMNode = node
+      this.handleRef(node)
     }
 
     // Append node to container if it isn't already
-    if (this.DOMNode.parentNode !== this.state.mountNode) {
+    if (this.DOMNode && this.DOMNode.parentNode !== this.state.mountNode) {
       if (insertAt === 'bottom') {
         this.state.mountNode.appendChild(this.DOMNode)
       } else {
@@ -179,7 +185,7 @@ class Portal extends Component<PortalProps, PortalState> {
       }
     }
 
-    return this.DOMNode
+    return this.DOMNode!
   }
 
   findMountNode(props: PortalProps) {