followup to Fix/textinput inputref

src/components/text-input/text-input.test.tsx

@@ -1,5 +1,6 @@
 import '@testing-library/jest-dom';
 import { render, screen } from '@testing-library/react';
+import { createRef } from 'react';
 import { InputType, TextInput } from './text-input';
 
 const testInputType = (type: InputType) => (): void => {
@@ -76,4 +77,125 @@ describe('<TextInput />', () => {
     const textInput = screen.getByTestId('textInput');
     expect(textInput).toHaveClass('a-text-input--full');
   });
+
+  it('attaches inputRef to the underlying input', () => {
+    const inputRef = createRef<HTMLInputElement>();
+    render(<TextInput id='ref-object' name='ref-object' inputRef={inputRef} />);
+
+    expect(inputRef.current).toBe(screen.getByTestId('textInput'));
+  });
+
+  it('supports callback inputRef', () => {
+    let node: HTMLInputElement | null = null;
+    render(
+      <TextInput
+        id='ref-callback'
+        name='ref-callback'
+        inputRef={(element) => {
+          node = element;
+        }}
+      />,
+    );
+
+    expect(node).toBe(screen.getByTestId('textInput'));
+  });
+
+  it('attaches both inputRef and forwarded ref', () => {
+    const inputRef = createRef<HTMLInputElement>();
+    const forwardedRef = createRef<HTMLInputElement>();
+    render(
+      <TextInput
+        id='both-refs'
+        name='both-refs'
+        inputRef={inputRef}
+        ref={forwardedRef}
+      />,
+    );
+
+    const textInput = screen.getByTestId('textInput');
+    expect(inputRef.current).toBe(textInput);
+    expect(forwardedRef.current).toBe(textInput);
+  });
+
+  it('keeps a stable inputRef attached across rerenders', () => {
+    const calls: (HTMLInputElement | null)[] = [];
+    const inputRef = (element: HTMLInputElement | null): void => {
+      calls.push(element);
+    };
+    const { rerender } = render(
+      <TextInput id='stable-ref' name='stable-ref' inputRef={inputRef} />,
+    );
+
+    rerender(
+      <TextInput id='stable-ref' name='stable-ref' inputRef={inputRef} />,
+    );
+
+    expect(calls).toEqual([screen.getByTestId('textInput')]);
+  });
+
+  it('clears object refs on unmount', () => {
+    const inputRef = createRef<HTMLInputElement>();
+    const forwardedRef = createRef<HTMLInputElement>();
+    const { unmount } = render(
+      <TextInput
+        id='unmount-ref'
+        name='unmount-ref'
+        inputRef={inputRef}
+        ref={forwardedRef}
+      />,
+    );
+
+    unmount();
+
+    expect(inputRef.current).toBeNull();
+    expect(forwardedRef.current).toBeNull();
+  });
+
+  it('runs ref cleanups on unmount', () => {
+    const calls: string[] = [];
+    const { unmount } = render(
+      <TextInput
+        id='cleanup-ref'
+        name='cleanup-ref'
+        inputRef={(element) => {
+          if (!element) return;
+          calls.push('input-attach');
+          return () => {
+            calls.push('input-cleanup');
+          };
+        }}
+        ref={(element) => {
+          if (!element) return;
+          calls.push('forwarded-attach');
+          return () => {
+            calls.push('forwarded-cleanup');
+          };
+        }}
+      />,
+    );
+
+    expect(calls).toEqual(['input-attach', 'forwarded-attach']);
+    unmount();
+    expect(calls).toEqual([
+      'input-attach',
+      'forwarded-attach',
+      'input-cleanup',
+      'forwarded-cleanup',
+    ]);
+  });
+
+  it('calls a callback inputRef with null on unmount when it returns no cleanup', () => {
+    const calls: (HTMLInputElement | null)[] = [];
+    const inputRef = (element: HTMLInputElement | null): void => {
+      calls.push(element);
+    };
+    const { unmount } = render(
+      <TextInput id='null-ref' name='null-ref' inputRef={inputRef} />,
+    );
+
+    unmount();
+
+    expect(calls.at(-1)).toBeNull();
+    expect(calls[0]).toBeInstanceOf(HTMLInputElement);
+  });
 });

src/components/text-input/text-input.tsx

@@ -2,14 +2,37 @@ import classnames from 'classnames';
 import {
   forwardRef,
   JSX,
+  useCallback,
   type ComponentPropsWithoutRef,
-  type ReactNode,
+  type Ref,
 } from 'react';
 import './cfpb-text-input.scss';
 import type { TextInputStatusType } from './text-input-status';
 import { getTextInputStatusClass } from './text-input-status';
 
-type TextInputReference = ReactNode;
+type TextInputReference = Ref<HTMLInputElement>;
+
+const releaseNothing = (): void => undefined;
+
+const bindReference = (
+  target: Ref<HTMLInputElement> | undefined,
+  element: HTMLInputElement | null,
+): (() => void) => {
+  if (typeof target === 'function') {
+    const cleanup = target(element);
+    return (): void => {
+      if (typeof cleanup === 'function') cleanup();
+      else target(null);
+    };
+  }
+  if (target) {
+    target.current = element;
+    return (): void => {
+      target.current = null;
+    };
+  }
+  return releaseNothing;
+};
 
 export type InputType =
   'email' | 'number' | 'password' | 'search' | 'tel' | 'text' | 'url';
@@ -42,7 +65,7 @@ export const TextInput = forwardRef<HTMLInputElement, TextInputProperties>(
     {
       className,
       id,
-      inputRef: _inputReference,
+      inputRef,
       isDisabled = false,
       name,
       status,
@@ -62,6 +85,18 @@ export const TextInput = forwardRef<HTMLInputElement, TextInputProperties>(
       classes.push('a-text-input--full');
     }
 
+    const setReferences = useCallback(
+      (element: HTMLInputElement | null): (() => void) => {
+        const releaseInputReference = bindReference(inputRef, element);
+        const releaseForwardedReference = bindReference(reference, element);
+        return (): void => {
+          releaseInputReference();
+          releaseForwardedReference();
+        };
+      },
+      [inputRef, reference],
+    );
+
     const input = (
       <input
         data-testid='textInput'
@@ -70,7 +105,7 @@ export const TextInput = forwardRef<HTMLInputElement, TextInputProperties>(
         id={id}
         name={name}
         type={type}
-        ref={reference}
+        ref={setReferences}
         {...otherInputProperties}
       />
     );