fix(mosaic): forward refs from Select.Group and Select.GroupLabel

Author: austincalvelage

packages/mosaic/src/primitives/select/__tests__/select.test.tsx

@@ -1,5 +1,6 @@
 import { cleanup, render, screen, waitFor } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
+import { createRef } from 'react';
 import { afterEach, describe, expect, it, vi } from 'vitest';
 
 import { axe } from '../../test-utils/axe';
@@ -993,6 +994,19 @@ describe('Select.Group', () => {
     expect(await axe(document.body, { rules: { region: { enabled: false } } })).toHaveNoViolations();
   });
 
+  it('forwards refs to the group and its label', () => {
+    const groupRef = createRef<HTMLDivElement>();
+    const labelRef = createRef<HTMLDivElement>();
+    render(
+      <Select.Group ref={groupRef}>
+        <Select.GroupLabel ref={labelRef}>Pome</Select.GroupLabel>
+      </Select.Group>,
+    );
+
+    expect(groupRef.current).toBe(screen.getByRole('group', { name: 'Pome' }));
+    expect(labelRef.current).toHaveTextContent('Pome');
+  });
+
   it('throws when the label is used outside <Select.Group>', () => {
     vi.spyOn(console, 'error').mockImplementation(() => {});
     expect(() => render(<Select.GroupLabel>Fruits</Select.GroupLabel>)).toThrow(/within <Select.Group>/);

packages/mosaic/src/primitives/select/select-group.tsx

@@ -1,6 +1,6 @@
 'use client';
 
-import { createContext, useContext, useId } from 'react';
+import React, { createContext, useContext, useId } from 'react';
 
 import { type ComponentProps, mergeProps, useRender } from '../utils';
 
@@ -9,7 +9,7 @@ const SelectGroupContext = createContext<string | undefined>(undefined);
 export type SelectGroupProps = ComponentProps<'div'>;
 
 /** A labelled set of options in the listbox, named by the `Select.GroupLabel` inside it. */
-export function SelectGroup(props: SelectGroupProps) {
+export const SelectGroup = React.forwardRef<HTMLDivElement, SelectGroupProps>(function SelectGroup(props, ref) {
   const { render, ...otherProps } = props;
   const labelId = useId();
 
@@ -23,25 +23,29 @@ export function SelectGroup(props: SelectGroupProps) {
       {useRender({
         defaultTagName: 'div',
         render,
+        ref,
         props: mergeProps<'div'>(defaultProps, otherProps),
       })}
     </SelectGroupContext.Provider>
   );
-}
+});
 
 export type SelectGroupLabelProps = ComponentProps<'div'>;
 
 /** Names the `Select.Group` it sits in. Throws outside one. */
-export function SelectGroupLabel(props: SelectGroupLabelProps) {
-  const { render, ...otherProps } = props;
-  const labelId = useContext(SelectGroupContext);
-  if (labelId === undefined) {
-    throw new Error('Select.GroupLabel must be used within <Select.Group>');
-  }
-
-  return useRender({
-    defaultTagName: 'div',
-    render,
-    props: mergeProps<'div'>({ id: labelId }, otherProps),
-  });
-}
+export const SelectGroupLabel = React.forwardRef<HTMLDivElement, SelectGroupLabelProps>(
+  function SelectGroupLabel(props, ref) {
+    const { render, ...otherProps } = props;
+    const labelId = useContext(SelectGroupContext);
+    if (labelId === undefined) {
+      throw new Error('Select.GroupLabel must be used within <Select.Group>');
+    }
+
+    return useRender({
+      defaultTagName: 'div',
+      render,
+      ref,
+      props: mergeProps<'div'>({ id: labelId }, otherProps),
+    });
+  },
+);