From 6cfe5e80d22bdb81bc8eea7fbca4818340d778ca Mon Sep 17 00:00:00 2001
From: Roy Johnson <roy.e.johnson@rice.edu>
Date: Thu, 21 Mar 2024 17:04:10 -0500
Subject: [PATCH] [DISCO-146] Fix colorpicker focus

---
 .../components/ColorPicker.spec.tsx           |  2 +-
 .../highlights/components/ColorPicker.tsx     | 22 ++++++++-----------
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/src/app/content/highlights/components/ColorPicker.spec.tsx b/src/app/content/highlights/components/ColorPicker.spec.tsx
index 3780cb5d16..9d7696ac2d 100644
--- a/src/app/content/highlights/components/ColorPicker.spec.tsx
+++ b/src/app/content/highlights/components/ColorPicker.spec.tsx
@@ -89,7 +89,7 @@ describe('ColorPicker', () => {
     const onRemove = jest.fn();
     const {root} = renderToDom(
       <TestContainer>
-        <ColorPicker onChange={onChange} onRemove={onRemove} />
+        <ColorPicker color={highlightStyles[0].label} onChange={onChange} onRemove={onRemove} />
       </TestContainer>
     );
     const rg = root.querySelector('[role="radiogroup"]') as HTMLDivElement;
diff --git a/src/app/content/highlights/components/ColorPicker.tsx b/src/app/content/highlights/components/ColorPicker.tsx
index c8a4a7d607..f961aaf092 100644
--- a/src/app/content/highlights/components/ColorPicker.tsx
+++ b/src/app/content/highlights/components/ColorPicker.tsx
@@ -104,22 +104,18 @@ const ColorPicker = ({className, ...props}: Props) => {
     },
     []
   );
-  const initialFocus = React.useCallback(
-    (event: React.FocusEvent) => {
-      if (!ref.current || event.currentTarget !== event.target) {
-        return;
-      }
-      const firstChecked = ref.current.querySelector<HTMLInputElement>('input[checked]');
-
-      if (firstChecked) {
-        firstChecked.focus();
-      } else {
-        ref.current.querySelector<HTMLInputElement>('input')?.focus();
+  const color = (props as SingleSelectProps).color;
+  const focusOnSelected = React.useCallback(
+    () => {
+      if (color && document?.activeElement === ref.current) {
+        ref.current?.querySelector<HTMLInputElement>(`input[name="${color}"]`)?.focus();
       }
     },
-    []
+    [color]
   );
 
+  React.useEffect(focusOnSelected, [focusOnSelected]);
+
   return (
     <div
       className={className}
@@ -128,7 +124,7 @@ const ColorPicker = ({className, ...props}: Props) => {
       aria-label='colors'
       ref={ref}
       onKeyDown={handleKeyNavigation}
-      onFocus={initialFocus}
+      onFocus={focusOnSelected}
     >
       {highlightStyles.map((style) => <ColorButton key={style.label}
         name={style.label}