From cdabe86df383a2403b56fd65d9743811ae7a08b4 Mon Sep 17 00:00:00 2001
From: Michael Winter <36558508+mcwinter07@users.noreply.github.com>
Date: Tue, 8 Oct 2024 14:20:18 +1100
Subject: [PATCH] Fix: future Select and FilterSelect page jump on focus issue
 (#5123)

* update overlay and listbox components so focus is manually set

* update future select portal story

* Add useIsClientReady hook

* update FilterSelect tests to increase functional coverage

* add a fallback focus and clean up code comments
---
 .changeset/lemon-tips-look.md                 |  5 +
 .../Filter/FilterSelect/FilterSelect.spec.tsx | 99 ++++++++++++++++++-
 .../src/Filter/FilterSelect/FilterSelect.tsx  |  3 +
 .../src/__future__/Select/Select.spec.tsx     |  8 +-
 .../src/__future__/Select/_docs/Select.mdx    |  4 +-
 .../Select/_docs/Select.stories.tsx           | 50 ++++++----
 .../Select/subcomponents/ListBox/ListBox.tsx  | 60 ++++++++++-
 .../Select/subcomponents/Overlay/Overlay.tsx  |  2 +-
 .../__utilities__/useIsClientReady/index.ts   |  1 +
 .../useIsClientReady/useIsClientReady.tsx     | 17 ++++
 10 files changed, 219 insertions(+), 30 deletions(-)
 create mode 100644 .changeset/lemon-tips-look.md
 create mode 100644 packages/components/src/__utilities__/useIsClientReady/index.ts
 create mode 100644 packages/components/src/__utilities__/useIsClientReady/useIsClientReady.tsx

diff --git a/.changeset/lemon-tips-look.md b/.changeset/lemon-tips-look.md
new file mode 100644
index 00000000000..aff10c614f6
--- /dev/null
+++ b/.changeset/lemon-tips-look.md
@@ -0,0 +1,5 @@
+---
+"@kaizen/components": patch
+---
+
+Fix for the tab focus jump issue on future Select and FilterSelect
diff --git a/packages/components/src/Filter/FilterSelect/FilterSelect.spec.tsx b/packages/components/src/Filter/FilterSelect/FilterSelect.spec.tsx
index f300b57bcdc..c04e092049b 100644
--- a/packages/components/src/Filter/FilterSelect/FilterSelect.spec.tsx
+++ b/packages/components/src/Filter/FilterSelect/FilterSelect.spec.tsx
@@ -48,7 +48,7 @@ describe("<FilterSelect>", () => {
   it("shows the options initially when isOpen is true", async () => {
     render(<FilterSelectWrapper isOpen />)
     await waitFor(() => {
-      expect(screen.queryByRole("listbox")).toBeVisible()
+      expect(screen.getByRole("listbox")).toBeVisible()
     })
   })
 
@@ -107,10 +107,82 @@ describe("<FilterSelect>", () => {
         render(<FilterSelectWrapper isOpen />)
         expect(screen.queryByRole("listbox")).toBeVisible()
         await waitFor(() => {
-          expect(screen.queryByRole("listbox")).toHaveFocus()
+          expect(screen.getAllByRole("option")[0]).toHaveFocus()
         })
       })
 
+      it("moves focus to the first item on ArrowDown if nothing has been selected", async () => {
+        render(<FilterSelectWrapper selectedKey={undefined} />)
+        const trigger = screen.getByRole("button", { name: "Coffee" })
+        await user.tab()
+        await waitFor(() => {
+          expect(trigger).toHaveFocus()
+        })
+        await user.keyboard("{ArrowDown}")
+
+        await waitFor(() => {
+          expect(screen.getAllByRole("option")[0]).toHaveFocus()
+        })
+      })
+      it("moves focus to the last item on ArrowUp if nothing has been selected", async () => {
+        render(<FilterSelectWrapper selectedKey={undefined} />)
+        const trigger = screen.getByRole("button", { name: "Coffee" })
+        await user.tab()
+        await waitFor(() => {
+          expect(trigger).toHaveFocus()
+        })
+        await user.keyboard("{ArrowUp}")
+
+        await waitFor(() => {
+          const options = screen.getAllByRole("option")
+          expect(options[options.length - 1]).toHaveFocus()
+        })
+      })
+      it("moves focus to the current selected item on Enter", async () => {
+        render(<FilterSelectWrapper selectedKey="hazelnut" />)
+        const trigger = screen.getByRole("button", {
+          name: "Coffee : Hazelnut",
+        })
+        await user.tab()
+        await waitFor(() => {
+          expect(trigger).toHaveFocus()
+        })
+        await user.keyboard("{Enter}")
+
+        await waitFor(() => {
+          expect(screen.getByRole("option", { name: "Hazelnut" })).toHaveFocus()
+        })
+      })
+      it("moves focus to the current selected item on ArrowUp", async () => {
+        render(<FilterSelectWrapper selectedKey="hazelnut" />)
+        const trigger = screen.getByRole("button", {
+          name: "Coffee : Hazelnut",
+        })
+        await user.tab()
+        await waitFor(() => {
+          expect(trigger).toHaveFocus()
+        })
+        await user.keyboard("{ArrowUp}")
+
+        await waitFor(() => {
+          expect(screen.getByRole("option", { name: "Hazelnut" })).toHaveFocus()
+        })
+      })
+      it("moves focus to the current selected item on ArrowDown", async () => {
+        render(<FilterSelectWrapper selectedKey="hazelnut" />)
+        const trigger = screen.getByRole("button", {
+          name: "Coffee : Hazelnut",
+        })
+        await user.tab()
+        await waitFor(() => {
+          expect(trigger).toHaveFocus()
+        })
+        await user.keyboard("{ArrowDown}")
+
+        await waitFor(() => {
+          expect(screen.getByRole("option", { name: "Hazelnut" })).toHaveFocus()
+        })
+      })
       it("closes when user hits the escape key", async () => {
         render(<FilterSelectWrapper isOpen />)
         expect(screen.queryByRole("listbox")).toBeVisible()
@@ -168,6 +240,29 @@ describe("<FilterSelect>", () => {
   })
 })
 
+describe("Stringified object values", () => {
+  it("finds selected option when value is a stringified object", () => {
+    const { getByRole } = render(
+      <FilterSelectWrapper
+        items={[
+          {
+            value: '{"sortBy":"creator_name","sortOrder":"asc"}',
+            label: "Created by A-Z",
+          },
+          {
+            value: '{"sortBy":"creator_name","sortOrder":"dsc"}',
+            label: "Created by Z-A",
+          },
+        ]}
+        selectedKey='{"sortBy":"creator_name","sortOrder":"asc"}'
+      />
+    )
+    expect(
+      getByRole("button", { name: "Coffee : Created by A-Z" })
+    ).toBeInTheDocument()
+  })
+})
+
 const defaultProps: FilterSelectProps = {
   label: "Coffee",
   isOpen: false,
diff --git a/packages/components/src/Filter/FilterSelect/FilterSelect.tsx b/packages/components/src/Filter/FilterSelect/FilterSelect.tsx
index 33020f0bc73..017e1c45fbe 100644
--- a/packages/components/src/Filter/FilterSelect/FilterSelect.tsx
+++ b/packages/components/src/Filter/FilterSelect/FilterSelect.tsx
@@ -78,6 +78,9 @@ export const FilterSelect = <Option extends SelectOption = SelectOption>({
   )
 
   const { buttonProps } = useButton(triggerProps, triggerRef)
+
+  // `aria-labelledby` and `aria-controls` are being remapped because the `buttonProps` ids generated by React Aria point to nothing.
+  // This should ideally be refactored but for now the `aria-controls` is set to the Filter's Listbox (menuProps.id) and the `aria-labelledby` to undefined so the accessible name is derived from the buttons content.
   const renderTriggerButtonProps = {
     ...buttonProps,
     "aria-labelledby": undefined,
diff --git a/packages/components/src/__future__/Select/Select.spec.tsx b/packages/components/src/__future__/Select/Select.spec.tsx
index 4008df83766..1964f8729b8 100644
--- a/packages/components/src/__future__/Select/Select.spec.tsx
+++ b/packages/components/src/__future__/Select/Select.spec.tsx
@@ -194,11 +194,13 @@ describe("<Select />", () => {
       })
 
       describe("Given the menu is opened", () => {
-        it("focuses the listbox initially", async () => {
-          const { getByRole } = render(<SelectWrapper defaultOpen />)
+        it("focuses on the first item", async () => {
+          const { getByRole, getAllByRole } = render(
+            <SelectWrapper defaultOpen />
+          )
           expect(getByRole("listbox")).toBeVisible()
           await waitFor(() => {
-            expect(getByRole("listbox")).toHaveFocus()
+            expect(getAllByRole("option")[0]).toHaveFocus()
           })
         })
         it("is closed when hits the escape key", async () => {
diff --git a/packages/components/src/__future__/Select/_docs/Select.mdx b/packages/components/src/__future__/Select/_docs/Select.mdx
index f540f6582a0..c1966567adc 100644
--- a/packages/components/src/__future__/Select/_docs/Select.mdx
+++ b/packages/components/src/__future__/Select/_docs/Select.mdx
@@ -98,8 +98,6 @@ Set `isFullWidth` to `true` to have the Select span the full width of its contai
 
 By default, the Select's popover will attach itself to the `body` of the document using React's `createPortal`.
 
-You can change the default behaviour by providing a `portalContainerId` to attach this to different element in the DOM. This can help to resolve issues that may arise with `z-index` or having a Select in a modal.
+You can change the default behavior by providing a `portalContainerId` to attach this to different element in the DOM. This can help to resolve issues that may arise with `z-index` or having a Select in a modal.
 
 <Canvas of={SelectStories.PortalContainer} />
-
-There is currently a known issue whereby a selected option will cause the page to scroll to the top of the window on open (click on [default to see example](https://cultureamp.design/?path=/docs/components-select-future--docs#portals)). This can be solved by setting a `portalContainerId` to the closest parent of the Select.
diff --git a/packages/components/src/__future__/Select/_docs/Select.stories.tsx b/packages/components/src/__future__/Select/_docs/Select.stories.tsx
index 45dcb170b58..9379d2368f6 100644
--- a/packages/components/src/__future__/Select/_docs/Select.stories.tsx
+++ b/packages/components/src/__future__/Select/_docs/Select.stories.tsx
@@ -1,5 +1,6 @@
 import React from "react"
 import { Meta, StoryObj } from "@storybook/react"
+import { ContextModal } from "~components/Modal"
 import { RadioField, RadioGroup } from "~components/Radio"
 import { Select } from "../Select"
 import { SelectOption } from "../types"
@@ -170,25 +171,40 @@ export const FullWidth: Story = {
 export const PortalContainer: Story = {
   render: args => {
     const portalContainerId = "id--portal-container"
+
+    const [isOpen, setIsOpen] = React.useState(false)
+
+    const handleOpen = (): void => setIsOpen(true)
+    const handleClose = (): void => setIsOpen(false)
     return (
       <>
-        <div
-          id={portalContainerId}
-          className="flex gap-24 bg-gray-200 p-12 overflow-hidden h-[200px] relative"
-        >
-          <Select
-            {...args}
-            label="Default"
-            selectedKey="batch-brew"
-            id="id--select-default"
-          />
-          <Select
-            {...args}
-            label="Inner portal"
-            selectedKey="batch-brew"
-            id="id--select-inner"
-            portalContainerId={portalContainerId}
-          />
+        <div className=" h-[500px] mb-24 block bg-gray-100 flex flex-col gap-16 justify-center items-center">
+          Page content
+          <button
+            type="button"
+            className="border border-gray-500"
+            onClick={handleOpen}
+          >
+            Open Modal
+          </button>
+          <ContextModal
+            isOpen={isOpen}
+            onConfirm={handleClose}
+            onDismiss={handleClose}
+            title="Select test"
+          >
+            <div
+              className="flex gap-24 bg-gray-200 p-12"
+              id={portalContainerId}
+            >
+              <Select
+                {...args}
+                label="Select within a modal"
+                id="id--select-inner"
+                portalContainerId={portalContainerId}
+              />
+            </div>
+          </ContextModal>
         </div>
       </>
     )
diff --git a/packages/components/src/__future__/Select/subcomponents/ListBox/ListBox.tsx b/packages/components/src/__future__/Select/subcomponents/ListBox/ListBox.tsx
index 79b8d090a6c..4c13f294a57 100644
--- a/packages/components/src/__future__/Select/subcomponents/ListBox/ListBox.tsx
+++ b/packages/components/src/__future__/Select/subcomponents/ListBox/ListBox.tsx
@@ -1,6 +1,8 @@
-import React, { HTMLAttributes } from "react"
+import React, { HTMLAttributes, Key, useEffect, useRef, ReactNode } from "react"
 import { AriaListBoxOptions, useListBox } from "@react-aria/listbox"
+import { SelectState } from "@react-stately/select"
 import classnames from "classnames"
+import { useIsClientReady } from "~components/__utilities__/useIsClientReady"
 import { OverrideClassName } from "~components/types/OverrideClassName"
 import { useSelectContext } from "../../context"
 import { SelectOption, SelectItem } from "../../types"
@@ -9,25 +11,75 @@ import styles from "./ListBox.module.scss"
 export type SingleListBoxProps<Option extends SelectOption> = OverrideClassName<
   HTMLAttributes<HTMLUListElement>
 > & {
-  children: React.ReactNode
+  children: ReactNode
   /** Props for the popup. */
   menuProps: AriaListBoxOptions<SelectItem<Option>>
 }
 
+/** A util to retrieve the key of the correct focusable items based of the focus strategy
+ * This is used to determine which element from the collection to focus to on open base on the keyboard event
+ * ie: UpArrow will set the focusStrategy to "last"
+ */
+const getOptionKeyFromCollection = (
+  state: SelectState<SelectItem<any>>
+): Key | null => {
+  if (state.selectedItem) {
+    return state.selectedItem.key
+  } else if (state.focusStrategy === "last") {
+    return state.collection.getLastKey()
+  }
+  return state.collection.getFirstKey()
+}
+
+/** This makes the use of query selector less brittle in instances where a failed selector is passed in
+ */
+const safeQuerySelector = (selector: string): HTMLElement | null => {
+  try {
+    return document.querySelector(selector)
+  } catch (error) {
+    // eslint-disable-next-line no-console
+    console.error("Kaizen querySelector failed:", error)
+    return null
+  }
+}
+
 export const ListBox = <Option extends SelectOption>({
   children,
   menuProps,
   classNameOverride,
   ...restProps
 }: SingleListBoxProps<Option>): JSX.Element => {
+  const isClientReady = useIsClientReady()
   const { state } = useSelectContext<Option>()
-  const ref = React.useRef<HTMLUListElement>(null)
+  const ref = useRef<HTMLUListElement>(null)
   const { listBoxProps } = useListBox(
-    { ...menuProps, disallowEmptySelection: true },
+    {
+      ...menuProps,
+      disallowEmptySelection: true,
+      // This is to ensure that the listbox doesn't use React Aria's auto focus feature for Listbox, which creates a visual bug
+      autoFocus: false,
+    },
     state,
     ref
   )
 
+  /**
+   * This uses the new useIsClientReady to ensure document exists before trying to querySelector and give the time to focus to the correct element
+   */
+  useEffect(() => {
+    if (isClientReady) {
+      const optionKey = getOptionKeyFromCollection(state)
+      const focusToElement = safeQuerySelector(`[data-key='${optionKey}']`)
+
+      if (focusToElement) {
+        focusToElement.focus()
+      } else {
+        // If an element is not found, focus on the listbox. This ensures the list can still be navigated to via keyboard if the keys do not align to the data attributes of the list items.
+        ref.current?.focus()
+      }
+    }
+  }, [isClientReady])
+
   return (
     <ul
       ref={ref}
diff --git a/packages/components/src/__future__/Select/subcomponents/Overlay/Overlay.tsx b/packages/components/src/__future__/Select/subcomponents/Overlay/Overlay.tsx
index 3088502dfcb..4d7082f463c 100644
--- a/packages/components/src/__future__/Select/subcomponents/Overlay/Overlay.tsx
+++ b/packages/components/src/__future__/Select/subcomponents/Overlay/Overlay.tsx
@@ -36,7 +36,7 @@ export const Overlay = <Option extends SelectOption>({
       {...restProps}
     >
       {/* eslint-disable-next-line jsx-a11y/no-autofocus */}
-      <FocusScope autoFocus restoreFocus>
+      <FocusScope autoFocus={false} restoreFocus>
         <DismissButton onDismiss={state.close} />
         {children}
         <DismissButton onDismiss={state.close} />
diff --git a/packages/components/src/__utilities__/useIsClientReady/index.ts b/packages/components/src/__utilities__/useIsClientReady/index.ts
new file mode 100644
index 00000000000..35ceae32e01
--- /dev/null
+++ b/packages/components/src/__utilities__/useIsClientReady/index.ts
@@ -0,0 +1 @@
+export * from "./useIsClientReady"
diff --git a/packages/components/src/__utilities__/useIsClientReady/useIsClientReady.tsx b/packages/components/src/__utilities__/useIsClientReady/useIsClientReady.tsx
new file mode 100644
index 00000000000..b3f2e9f276f
--- /dev/null
+++ b/packages/components/src/__utilities__/useIsClientReady/useIsClientReady.tsx
@@ -0,0 +1,17 @@
+import { useState, useEffect } from "react"
+
+/**
+ * A hook that returns a truthy value indicating if the code can be run on client side.
+ * This is a useful hook for determining if the `document` or `window` objects are available.
+ */
+export const useIsClientReady = (): boolean => {
+  const [isClientReady, setIsClientReady] = useState(false)
+
+  useEffect(() => {
+    if (typeof window !== "undefined" && typeof document !== "undefined") {
+      setIsClientReady(true)
+    }
+  }, [])
+
+  return isClientReady
+}