TextInputWithTokens bugfix: focus next token when removing first toke…

…n with keyboard (#1529)

* fixes a bug when removing the first token with the keyboard

Before this fix, the next token would not get focused after removing the first token in the input by selecting it and hitting the 'Backspace' or 'Delete' key

* adds changeset

* Update .changeset/eighty-beds-sneeze.md

Co-authored-by: Cole Bemis <colebemis@github.com>

* Update .changeset/eighty-beds-sneeze.md

Co-authored-by: Cole Bemis <colebemis@github.com>

Co-authored-by: Cole Bemis <colebemis@github.com>

.changeset/eighty-beds-sneeze.md

@@ -0,0 +1,5 @@
+---
+'@primer/components': patch
+---
+
+Fixes a bug in `TextInputWithTokens` where the next token would not receive focus after the user deleted the first token using the keyboard

src/TextInputWithTokens.tsx

@@ -11,6 +11,7 @@ import {useProvidedRefOrCreate} from './hooks'
 import UnstyledTextInput from './_UnstyledTextInput'
 import TextInputWrapper from './_TextInputWrapper'
 import Box from './Box'
+import {isFocusable} from './utils/iterateFocusableElements'
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 type AnyReactComponent = React.ComponentType<any>
@@ -117,10 +118,24 @@ function TextInputWithTokensInnerComponent<TokenComponentType extends AnyReactCo
   const handleTokenRemove = (tokenId: string | number) => {
     onTokenRemove(tokenId)
 
-    if (selectedTokenIndex) {
-      const nextElementToFocus = containerRef.current?.children[selectedTokenIndex] as HTMLElement
-      nextElementToFocus.focus()
-    }
+    // HACK: wait a tick for the the token node to be removed from the DOM
+    setTimeout(() => {
+      const nextElementToFocus = containerRef.current?.children[selectedTokenIndex || 0] as HTMLElement | undefined
+
+      // when removing the first token by keying "Backspace" or "Delete",
+      // `nextFocusableElement` is the div that wraps the input
+      const firstFocusable =
+        nextElementToFocus && isFocusable(nextElementToFocus)
+          ? nextElementToFocus
+          : (Array.from(containerRef.current?.children || []) as HTMLElement[]).find(el => isFocusable(el))
+
+      if (firstFocusable) {
+        firstFocusable.focus()
+      } else {
+        // if there are no tokens left, focus the input
+        ref.current?.focus()
+      }
+    }, 0)
   }
 
   const handleTokenFocus: (tokenIndex: number) => FocusEventHandler = tokenIndex => () => {

src/__tests__/TextInputWithTokens.test.tsx

@@ -238,6 +238,44 @@ describe('TextInputWithTokens', () => {
     expect(onRemoveMock).toHaveBeenCalledWith(mockTokens[4].id)
   })
 
+  it('moves focus to the next token when removing the first token', () => {
+    jest.useFakeTimers()
+    const onRemoveMock = jest.fn()
+    const {getByText} = HTMLRender(
+      <TextInputWithTokens tokens={[...mockTokens].slice(0, 2)} onTokenRemove={onRemoveMock} />
+    )
+    const tokenNode = getByText(mockTokens[0].text)
+
+    fireEvent.focus(tokenNode)
+    fireEvent.keyDown(tokenNode, {key: 'Backspace'})
+
+    jest.runAllTimers()
+    setTimeout(() => {
+      expect(document.activeElement?.textContent).toBe(mockTokens[1].text)
+    }, 0)
+
+    jest.useRealTimers()
+  })
+
+  it('moves focus to the input when the last token is removed', () => {
+    jest.useFakeTimers()
+    const onRemoveMock = jest.fn()
+    const {getByText, getByLabelText} = HTMLRender(
+      <LabelledTextInputWithTokens tokens={[mockTokens[0]]} onTokenRemove={onRemoveMock} />
+    )
+    const tokenNode = getByText(mockTokens[0].text)
+    const inputNode = getByLabelText('Tokens')
+
+    fireEvent.focus(tokenNode)
+    fireEvent.keyDown(tokenNode, {key: 'Backspace'})
+
+    jest.runAllTimers()
+    setTimeout(() => {
+      expect(document.activeElement?.id).toBe(inputNode.id)
+    }, 0)
+    jest.useRealTimers()
+  })
+
   it('calls onKeyDown', () => {
     const onRemoveMock = jest.fn()
     const onKeyDownMock = jest.fn()

src/stories/TextInputWithTokens.stories.tsx

@@ -47,7 +47,7 @@ const mockTokens = [
 ]
 
 export const Default = () => {
-  const [tokens, setTokens] = useState([...mockTokens].slice(0, 2))
+  const [tokens, setTokens] = useState([...mockTokens].slice(0, 3))
   const onTokenRemove: (tokenId: string | number) => void = tokenId => {
     setTokens(tokens.filter(token => token.id !== tokenId))
   }