From 81a4c8039b7e42d974d5e96ab5bb0fb9098866fa Mon Sep 17 00:00:00 2001
From: Satyapriya Bai <v-satybai@microsoft.com>
Date: Wed, 19 Jun 2024 15:12:44 +0530
Subject: [PATCH] [accessibility-2262594]:[Keyboard Navigation - Azure Cosmos
 DB - Data Explorer]: Keyboard focus is not retaining back to 'more' button
 after closing 'Delete container' dialog.

---
 src/Explorer/ContextMenuButtonFactory.tsx          | 14 ++++++++++----
 .../Controls/TreeComponent/TreeNodeComponent.tsx   |  8 +++++---
 .../__snapshots__/TreeNodeComponent.test.tsx.snap  |  8 ++++----
 .../DeleteCollectionConfirmationPane.tsx           | 14 ++++++++++++--
 .../Panes/DeleteDatabaseConfirmationPanel.tsx      | 13 +++++++++++--
 5 files changed, 42 insertions(+), 15 deletions(-)

diff --git a/src/Explorer/ContextMenuButtonFactory.tsx b/src/Explorer/ContextMenuButtonFactory.tsx
index 90c78e365..d4186ba27 100644
--- a/src/Explorer/ContextMenuButtonFactory.tsx
+++ b/src/Explorer/ContextMenuButtonFactory.tsx
@@ -52,12 +52,15 @@ export const createDatabaseContextMenu = (container: Explorer, databaseId: strin
   if (userContext.apiType !== "Tables" || userContext.features.enableSDKoperations) {
     items.push({
       iconSrc: DeleteDatabaseIcon,
-      onClick: () =>
+      onClick: (lastFocusedElement?: React.RefObject<HTMLElement>) =>
         useSidePanel
           .getState()
           .openSidePanel(
             "Delete " + getDatabaseName(),
-            <DeleteDatabaseConfirmationPanel refreshDatabases={() => container.refreshAllDatabases()} />,
+            <DeleteDatabaseConfirmationPanel
+              lastFocusedElement={lastFocusedElement}
+              refreshDatabases={() => container.refreshAllDatabases()}
+            />,
           ),
       label: `Delete ${getDatabaseName()}`,
       styleClass: "deleteDatabaseMenuItem",
@@ -132,13 +135,16 @@ export const createCollectionContextMenuButton = (
   if (configContext.platform !== Platform.Fabric) {
     items.push({
       iconSrc: DeleteCollectionIcon,
-      onClick: () => {
+      onClick: (lastFocusedElement?: React.RefObject<HTMLElement>) => {
         useSelectedNode.getState().setSelectedNode(selectedCollection);
         useSidePanel
           .getState()
           .openSidePanel(
             "Delete " + getCollectionName(),
-            <DeleteCollectionConfirmationPane refreshDatabases={() => container.refreshAllDatabases()} />,
+            <DeleteCollectionConfirmationPane
+              lastFocusedElement={lastFocusedElement}
+              refreshDatabases={() => container.refreshAllDatabases()}
+            />,
           );
       },
       label: `Delete ${getCollectionName()}`,
diff --git a/src/Explorer/Controls/TreeComponent/TreeNodeComponent.tsx b/src/Explorer/Controls/TreeComponent/TreeNodeComponent.tsx
index 9ec2c904a..352b732b8 100644
--- a/src/Explorer/Controls/TreeComponent/TreeNodeComponent.tsx
+++ b/src/Explorer/Controls/TreeComponent/TreeNodeComponent.tsx
@@ -21,7 +21,7 @@ import { useCallback } from "react";
 
 export interface TreeNodeMenuItem {
   label: string;
-  onClick: () => void;
+  onClick: (value?: React.RefObject<HTMLElement>) => void;
   iconSrc?: string;
   isDisabled?: boolean;
   styleClass?: string;
@@ -73,6 +73,7 @@ export const TreeNodeComponent: React.FC<TreeNodeComponentProps> = ({
   treeNodeId,
 }: TreeNodeComponentProps): JSX.Element => {
   const [isLoading, setIsLoading] = React.useState<boolean>(false);
+  const contextMenuRef = React.useRef<HTMLButtonElement>(null);
 
   const getSortedChildren = (treeNode: TreeNode): TreeNode[] => {
     if (!treeNode || !treeNode.children) {
@@ -139,7 +140,7 @@ export const TreeNodeComponent: React.FC<TreeNodeComponentProps> = ({
       data-test={`TreeNode/ContextMenuItem:${menuItem.label}`}
       disabled={menuItem.isDisabled}
       key={menuItem.label}
-      onClick={menuItem.onClick}
+      onClick={() => menuItem.onClick(contextMenuRef)}
     >
       {menuItem.label}
     </MenuItem>
@@ -158,11 +159,12 @@ export const TreeNodeComponent: React.FC<TreeNodeComponentProps> = ({
         actions={
           contextMenuItems.length > 0 && (
             <Menu onOpenChange={onMenuOpenChange}>
-              <MenuTrigger disableButtonEnhancement>
+              <MenuTrigger disableButtonEnhancement={true}>
                 <Button
                   aria-label="More options"
                   data-test="TreeNode/ContextMenuTrigger"
                   appearance="subtle"
+                  ref={contextMenuRef}
                   icon={<MoreHorizontal20Regular />}
                 />
               </MenuTrigger>
diff --git a/src/Explorer/Controls/TreeComponent/__snapshots__/TreeNodeComponent.test.tsx.snap b/src/Explorer/Controls/TreeComponent/__snapshots__/TreeNodeComponent.test.tsx.snap
index 658e01e0c..69697883d 100644
--- a/src/Explorer/Controls/TreeComponent/__snapshots__/TreeNodeComponent.test.tsx.snap
+++ b/src/Explorer/Controls/TreeComponent/__snapshots__/TreeNodeComponent.test.tsx.snap
@@ -1560,14 +1560,14 @@ exports[`TreeNodeComponent renders a node with a menu 1`] = `
               <MenuList>
                 <MenuItem
                   data-test="TreeNode/ContextMenuItem:enabledItem"
-                  onClick={[MockFunction enabledItemClick]}
+                  onClick={[Function]}
                 >
                   enabledItem
                 </MenuItem>
                 <MenuItem
                   data-test="TreeNode/ContextMenuItem:disabledItem"
                   disabled={true}
-                  onClick={[MockFunction disabledItemClick]}
+                  onClick={[Function]}
                 >
                   disabledItem
                 </MenuItem>
@@ -1604,7 +1604,7 @@ exports[`TreeNodeComponent renders a node with a menu 1`] = `
       <MenuItem
         data-test="TreeNode/ContextMenuItem:enabledItem"
         key="enabledItem"
-        onClick={[MockFunction enabledItemClick]}
+        onClick={[Function]}
       >
         enabledItem
       </MenuItem>
@@ -1612,7 +1612,7 @@ exports[`TreeNodeComponent renders a node with a menu 1`] = `
         data-test="TreeNode/ContextMenuItem:disabledItem"
         disabled={true}
         key="disabledItem"
-        onClick={[MockFunction disabledItemClick]}
+        onClick={[Function]}
       >
         disabledItem
       </MenuItem>
diff --git a/src/Explorer/Panes/DeleteCollectionConfirmationPane/DeleteCollectionConfirmationPane.tsx b/src/Explorer/Panes/DeleteCollectionConfirmationPane/DeleteCollectionConfirmationPane.tsx
index 05eff191f..581784739 100644
--- a/src/Explorer/Panes/DeleteCollectionConfirmationPane/DeleteCollectionConfirmationPane.tsx
+++ b/src/Explorer/Panes/DeleteCollectionConfirmationPane/DeleteCollectionConfirmationPane.tsx
@@ -12,17 +12,19 @@ import { getCollectionName } from "Utils/APITypeUtils";
 import * as NotificationConsoleUtils from "Utils/NotificationConsoleUtils";
 import { useSidePanel } from "hooks/useSidePanel";
 import { useTabs } from "hooks/useTabs";
-import React, { FunctionComponent, useState } from "react";
+import React, { FunctionComponent, useEffect, useState } from "react";
 import { useDatabases } from "../../useDatabases";
 import { useSelectedNode } from "../../useSelectedNode";
 import { RightPaneForm, RightPaneFormProps } from "../RightPaneForm/RightPaneForm";
 
 export interface DeleteCollectionConfirmationPaneProps {
   refreshDatabases: () => Promise<void>;
+  lastFocusedElement: React.MutableRefObject<HTMLElement>;
 }
 
 export const DeleteCollectionConfirmationPane: FunctionComponent<DeleteCollectionConfirmationPaneProps> = ({
   refreshDatabases,
+  lastFocusedElement,
 }: DeleteCollectionConfirmationPaneProps) => {
   const closeSidePanel = useSidePanel((state) => state.closeSidePanel);
   const [deleteCollectionFeedback, setDeleteCollectionFeedback] = useState<string>("");
@@ -35,7 +37,7 @@ export const DeleteCollectionConfirmationPane: FunctionComponent<DeleteCollectio
 
   const collectionName = getCollectionName().toLocaleLowerCase();
   const paneTitle = "Delete " + collectionName;
-
+  const lastItemElement = lastFocusedElement?.current;
   const onSubmit = async (): Promise<void> => {
     const collection = useSelectedNode.getState().findSelectedCollection();
     if (!collection || inputCollectionName !== collection.id()) {
@@ -111,6 +113,14 @@ export const DeleteCollectionConfirmationPane: FunctionComponent<DeleteCollectio
   };
   const confirmContainer = `Confirm by typing the ${collectionName.toLowerCase()} id`;
   const reasonInfo = `Help us improve Azure Cosmos DB! What is the reason why you are deleting this ${collectionName}?`;
+  useEffect(() => {
+    return () => {
+      if (lastItemElement) {
+        lastItemElement.focus();
+      }
+    };
+  }, [lastItemElement]);
+
   return (
     <RightPaneForm {...props}>
       <div className="panelFormWrapper">
diff --git a/src/Explorer/Panes/DeleteDatabaseConfirmationPanel.tsx b/src/Explorer/Panes/DeleteDatabaseConfirmationPanel.tsx
index af291bba6..a228462fa 100644
--- a/src/Explorer/Panes/DeleteDatabaseConfirmationPanel.tsx
+++ b/src/Explorer/Panes/DeleteDatabaseConfirmationPanel.tsx
@@ -13,7 +13,7 @@ import { getDatabaseName } from "Utils/APITypeUtils";
 import { logConsoleError } from "Utils/NotificationConsoleUtils";
 import { useSidePanel } from "hooks/useSidePanel";
 import { useTabs } from "hooks/useTabs";
-import React, { FunctionComponent, useState } from "react";
+import React, { FunctionComponent, useEffect, useState } from "react";
 import { useDatabases } from "../useDatabases";
 import { useSelectedNode } from "../useSelectedNode";
 import { PanelInfoErrorComponent, PanelInfoErrorProps } from "./PanelInfoErrorComponent";
@@ -21,10 +21,12 @@ import { RightPaneForm, RightPaneFormProps } from "./RightPaneForm/RightPaneForm
 
 interface DeleteDatabaseConfirmationPanelProps {
   refreshDatabases: () => Promise<void>;
+  lastFocusedElement: React.MutableRefObject<HTMLElement>;
 }
 
 export const DeleteDatabaseConfirmationPanel: FunctionComponent<DeleteDatabaseConfirmationPanelProps> = ({
   refreshDatabases,
+  lastFocusedElement,
 }: DeleteDatabaseConfirmationPanelProps): JSX.Element => {
   const closeSidePanel = useSidePanel((state) => state.closeSidePanel);
   const isLastNonEmptyDatabase = useDatabases((state) => state.isLastNonEmptyDatabase);
@@ -34,7 +36,7 @@ export const DeleteDatabaseConfirmationPanel: FunctionComponent<DeleteDatabaseCo
   const [databaseInput, setDatabaseInput] = useState<string>("");
   const [databaseFeedbackInput, setDatabaseFeedbackInput] = useState<string>("");
   const selectedDatabase: Database = useDatabases.getState().findSelectedDatabase();
-
+  const lastItemElement = lastFocusedElement?.current;
   const submit = async (): Promise<void> => {
     if (selectedDatabase?.id() && databaseInput !== selectedDatabase.id()) {
       setFormError(
@@ -126,6 +128,13 @@ export const DeleteDatabaseConfirmationPanel: FunctionComponent<DeleteDatabaseCo
   };
   const confirmDatabase = `Confirm by typing the ${getDatabaseName()} id`;
   const reasonInfo = `Help us improve Azure Cosmos DB! What is the reason why you are deleting this ${getDatabaseName()}?`;
+  useEffect(() => {
+    return () => {
+      if (lastItemElement) {
+        lastItemElement.focus();
+      }
+    };
+  }, [lastItemElement]);
   return (
     <RightPaneForm {...props}>
       {!formError && <PanelInfoErrorComponent {...errorProps} />}