update useProvidedRefOrCreate hook to support function refs (#1081)

Author: rezrah

.changeset/sixty-walls-pull.md

@@ -0,0 +1,7 @@
+---
+'@primer/react-brand': patch
+---
+
+Updated `useProvidedRefOrCreate` to support functional refs. Backwards compatibility with `RefObject` remains unchanged.
+
+This update improves support for functional refs in the following components: `Accordion`, `Checkbox`, `Radio`, `SubNav`, `Tooltip` and `VideoPlayer`.

packages/react/src/forms/Checkbox/Checkbox.test.tsx

@@ -0,0 +1,106 @@
+import React from 'react'
+import {render, cleanup} from '@testing-library/react'
+import userEvent from '@testing-library/user-event'
+import '@testing-library/jest-dom'
+import {toHaveNoViolations} from 'jest-axe'
+
+import {Checkbox} from '../Checkbox'
+
+expect.extend(toHaveNoViolations)
+
+describe('Checkbox', () => {
+  afterEach(cleanup)
+
+  it('renders a checkbox correctly into the document', () => {
+    const {getByRole} = render(<Checkbox value="one" defaultChecked />)
+
+    const checkbox = getByRole('checkbox')
+    expect(checkbox).toBeInTheDocument()
+    expect(checkbox).toBeChecked()
+    expect(checkbox).toHaveAttribute('value', 'one')
+  })
+
+  it('toggles the checkbox when clicked', async () => {
+    const user = userEvent.setup()
+
+    const {getByRole} = render(<Checkbox value="one" defaultChecked />)
+
+    const checkbox = getByRole('checkbox')
+
+    expect(checkbox).toBeChecked()
+
+    await user.click(checkbox)
+    expect(checkbox).not.toBeChecked()
+
+    await user.click(checkbox)
+    expect(checkbox).toBeChecked()
+  })
+
+  it('toggles the checkbox when the associated span is clicked', async () => {
+    const user = userEvent.setup()
+
+    const {getByRole, container} = render(<Checkbox value="one" defaultChecked />)
+
+    const checkbox = getByRole('checkbox')
+    const span = container.querySelector('span.Checkbox') as HTMLSpanElement
+
+    expect(span).toBeInTheDocument()
+    expect(checkbox).toBeChecked()
+
+    await user.click(span)
+    expect(checkbox).not.toBeChecked()
+
+    await user.click(span)
+    expect(checkbox).toBeChecked()
+  })
+
+  it('toggles the checkbox when the associated span is clicked and a functional ref is passed', async () => {
+    const user = userEvent.setup()
+    const mockRef = jest.fn()
+
+    const {getByRole, container} = render(<Checkbox value="one" defaultChecked ref={mockRef} />)
+
+    const checkbox = getByRole('checkbox')
+    const span = container.querySelector('span.Checkbox') as HTMLSpanElement
+
+    expect(mockRef).toHaveBeenCalledWith(checkbox)
+    expect(checkbox).toBeChecked()
+
+    await user.click(span)
+    expect(checkbox).not.toBeChecked()
+
+    await user.click(span)
+    expect(checkbox).toBeChecked()
+  })
+
+  it('remains clickable when passing an arbitrary function as a ref', async () => {
+    // eslint-disable-next-line @typescript-eslint/no-empty-function
+    const {getByRole} = render(<Checkbox value="test" ref={() => {}} />)
+
+    const checkbox = getByRole('checkbox')
+    expect(checkbox).not.toBeChecked()
+
+    const user = userEvent.setup()
+    await user.click(checkbox)
+    expect(checkbox).toBeChecked()
+  })
+
+  it('supports a standard RefObject', async () => {
+    const user = userEvent.setup()
+    const refObject = React.createRef<HTMLInputElement>()
+
+    const {getByRole} = render(<Checkbox value="test" ref={refObject} />)
+
+    const checkbox = getByRole('checkbox')
+
+    expect(refObject.current).toBe(checkbox)
+    expect(checkbox).not.toBeChecked()
+
+    if (refObject.current) {
+      await user.click(refObject.current)
+    }
+
+    expect(checkbox).toBeChecked()
+    expect(refObject.current?.checked).toBe(true)
+  })
+})

packages/react/src/hooks/useRef.ts

@@ -1,16 +1,52 @@
-import React, {useRef} from 'react'
+import React, {useRef, useMemo} from 'react'
 
 /**
  * There are some situations where we only want to create a new ref if one is not provided to a component
  * or hook as a prop. However, due to the `rules-of-hooks`, we cannot conditionally make a call to `React.useRef`
  * only in the situations where the ref is not provided as a prop.
  * This hook aims to encapsulate that logic, so the consumer doesn't need to be concerned with violating `rules-of-hooks`.
- * @param providedRef The ref to use - if undefined, will use the ref from a call to React.useRef
+ * @param providedRef The ref to use - can be a RefObject, RefCallback (functional ref), null or undefined.
  * @type TRef The type of the RefObject which should be created.
+ * @returns A RefObject that either uses the provided ref, wraps a functional ref, or creates a new one
  */
-
-export function useProvidedRefOrCreate<TRef>(providedRef?: React.RefObject<TRef>): React.RefObject<TRef> {
+export function useProvidedRefOrCreate<TRef>(
+  providedRef?: React.RefObject<TRef> | React.RefCallback<TRef> | null,
+): React.RefObject<TRef> {
   const createdRef = useRef<TRef>(null)
+  const wrapperRef = useRef<React.MutableRefObject<TRef | null> | null>(null)
+
+  const finalRef = useMemo(() => {
+    // Handle function refs
+    // We wrap it in a MutableRefObject to ensure we can set the current value
+    // and also call the callback with the latest value
+    if (typeof providedRef === 'function') {
+      if (!wrapperRef.current) {
+        let currentVal: TRef | null = null
+
+        wrapperRef.current = {
+          get current() {
+            return currentVal
+          },
+          set current(value: TRef | null) {
+            currentVal = value
+
+            providedRef(value)
+          },
+        }
+      }
+      return wrapperRef.current
+    }
+
+    wrapperRef.current = null
+
+    // Return the provided ref if it's a RefObject
+    if (providedRef && typeof providedRef === 'object' && providedRef.current !== undefined) {
+      return providedRef
+    }
+
+    // Return a new ref if it's neither a function nor a RefObject
+    return createdRef
+  }, [providedRef, createdRef])
 
-  return providedRef ?? createdRef
+  return finalRef
 }