From da566044649a114e23789e3ac90c022cdbf5eaeb Mon Sep 17 00:00:00 2001
From: Mike Perrotti <mperrotti@github.com>
Date: Wed, 27 Oct 2021 14:51:21 -0400
Subject: [PATCH] TextInputWithTokens bugfix: focus next token when removing
 first token 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            |  5 +++
 src/TextInputWithTokens.tsx                 | 23 ++++++++++---
 src/__tests__/TextInputWithTokens.test.tsx  | 38 +++++++++++++++++++++
 src/stories/TextInputWithTokens.stories.tsx |  2 +-
 4 files changed, 63 insertions(+), 5 deletions(-)
 create mode 100644 .changeset/eighty-beds-sneeze.md

diff --git a/.changeset/eighty-beds-sneeze.md b/.changeset/eighty-beds-sneeze.md
new file mode 100644
index 00000000000..5391487e71c
--- /dev/null
+++ b/.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
diff --git a/src/TextInputWithTokens.tsx b/src/TextInputWithTokens.tsx
index 7868874d66a..f4d50260077 100644
--- a/src/TextInputWithTokens.tsx
+++ b/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 => () => {
diff --git a/src/__tests__/TextInputWithTokens.test.tsx b/src/__tests__/TextInputWithTokens.test.tsx
index 2669bb7c91d..8f44563ed43 100644
--- a/src/__tests__/TextInputWithTokens.test.tsx
+++ b/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()
diff --git a/src/stories/TextInputWithTokens.stories.tsx b/src/stories/TextInputWithTokens.stories.tsx
index 24da3697499..494ed870144 100644
--- a/src/stories/TextInputWithTokens.stories.tsx
+++ b/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))
   }