Better docs and correctly handle override IDs changing

src/utilities/unique-id/hooks.ts

@@ -1,26 +1,38 @@
 import {useContext, useRef} from 'react';
 import {UniqueIdFactoryContext} from './context';
 
-export function useUniqueId(prefix = '', override?: string) {
+/**
+ * Returns a unique id that remains consistent across multiple rerenders of the
+ * same hook
+ * @param prefix Defines a prefix for the ID. You probably want to set this to
+ *   the name of the component you're calling `useUniqueId` in.
+ * @param overrideId Defines a fixed value to use instead of generating a unique
+ *   ID. Useful for components that allow consumers to specify a fixed ID.
+ */
+export function useUniqueId(prefix = '', overrideId = '') {
   const idFactory = useContext(UniqueIdFactoryContext);
 
+  // By using a ref to store the uniqueId for each incovation of the hook and
+  // checking that it is not already populated we ensure that we don't generate
+  // a new ID on ever rerender of a component.
+  const uniqueIdRef = useRef<string | null>(null);
+
   if (!idFactory) {
     throw new Error(
       'No UniqueIdFactory was provided. Your application must be wrapped in an <AppProvider> component. See https://polaris.shopify.com/components/structure/app-provider for implementation instructions.',
     );
   }
 
-  // Use a ref to ensure that the returned ID does not increment on every
-  // rerendering of a component
-  // The first time a component is rendered the ref will be empty.
-  // In that case fill it with the next available ID
-  // On subsequent renders we use the existing value instead of using a new id
-  const currentComponentIdRef = useRef<string | null>(null);
+  // If an override was specified, then use that instead of using a unique ID
+  // Hooks can't be called conditionally so this has to go after all use* calls
+  if (overrideId) {
+    return overrideId;
+  }
 
-  if (!currentComponentIdRef.current) {
-    // If an override was specified, then use that instead of incrementing the ID
-    currentComponentIdRef.current = override || idFactory.nextId(prefix);
+  // If a unique id has not yet been generated, then get a new one
+  if (!uniqueIdRef.current) {
+    uniqueIdRef.current = idFactory.nextId(prefix);
   }
 
-  return currentComponentIdRef.current;
+  return uniqueIdRef.current;
 }

src/utilities/unique-id/tests/hooks.test.tsx

@@ -130,4 +130,53 @@ describe('useUniqueId', () => {
 
     expect(harness.find(HasProp)!.html()).toStrictEqual('count2 :: polaris-1');
   });
+
+  it.only('updates the ID if the overridden ID changes', () => {
+    type HasPropProps = {info: string; idOverride?: string};
+    const HasProp = ({info, idOverride}: HasPropProps) => (
+      <React.Fragment>
+        {info} :: {useUniqueId('', idOverride)}
+      </React.Fragment>
+    );
+
+    const ReRenderingTestHarness = () => {
+      const [count, setCount] = React.useState(1);
+      const incrementCount = React.useCallback(
+        () => setCount((count) => count + 1),
+        [],
+      );
+
+      // eslint-disable-next-line shopify/jest/no-if
+      const override = count % 2 === 0 ? `override-${count}` : undefined;
+
+      return (
+        <React.Fragment>
+          <button onClick={incrementCount}>Click Me</button>
+          <HasProp info={`count${count}`} idOverride={override} />
+        </React.Fragment>
+      );
+    };
+
+    const harness = mountWithApp(<ReRenderingTestHarness />);
+
+    // Initially we use an incremental id
+    expect(harness.find(HasProp)!.html()).toStrictEqual('count1 :: polaris-1');
+
+    // But then we set an override id, so it should use that
+    harness.find('button')!.trigger('onClick');
+    expect(harness.find(HasProp)!.html()).toStrictEqual('count2 :: override-2');
+
+    // Then on the next render we don't set an override id, so we should go back
+    // to using the incremental id
+    harness.find('button')!.trigger('onClick');
+    expect(harness.find(HasProp)!.html()).toStrictEqual('count3 :: polaris-1');
+
+    // Back to setting an override id
+    harness.find('button')!.trigger('onClick');
+    expect(harness.find(HasProp)!.html()).toStrictEqual('count4 :: override-4');
+
+    // Back to not setting an override, so back to using the incremental id
+    harness.find('button')!.trigger('onClick');
+    expect(harness.find(HasProp)!.html()).toStrictEqual('count5 :: polaris-1');
+  });
 });