[fix] Add a wait condition/state for popover focus

- this makes it so EuiContextMenu doesn't call `.focus()` too early and hijack the `returnFocus` element that our focus trap dependency sets via `document.activeElement` - see #5760 (comment)

+ Add Cypress E2E tests for popover close focus return (w/ bonus `initialFocus` regression test)

src/components/context_menu/context_menu_panel.spec.tsx

@@ -8,7 +8,7 @@
 
 /// <reference types="../../../cypress/support"/>
 
-import React from 'react';
+import React, { useState } from 'react';
 
 import { EuiPopover } from '../popover';
 import { EuiContextMenu } from './context_menu';
@@ -123,17 +123,69 @@ describe('EuiContextMenuPanel', () => {
       });
     });
 
-    describe('when inside an EuiPopover', () => {
-      it('reclaims focus from the parent popover panel', () => {
-        cy.mount(
-          <EuiPopover isOpen={true} button={<button />}>
-            <EuiContextMenuPanel items={items} />
+    describe('within an EuiPopover', () => {
+      const ContextMenuInPopover: React.FC<any> = ({ children, ...rest }) => {
+        const [isOpen, setIsOpen] = useState(false);
+        const closePopover = () => setIsOpen(false);
+        const openPopover = () => setIsOpen(true);
+        return (
+          <EuiPopover
+            isOpen={isOpen}
+            closePopover={closePopover}
+            button={
+              <button data-test-subj="popoverToggle" onClick={openPopover}>
+                Toggle popover
+              </button>
+            }
+            {...rest}
+          >
+            <EuiContextMenuPanel
+              items={[
+                children || (
+                  <button onClick={closePopover}>
+                    Closes popover from context menu
+                  </button>
+                ),
+              ]}
+            />
           </EuiPopover>
         );
-        cy.wait(400); // EuiPopover's updateFocus() takes ~350ms to run
+      };
+
+      const mountAndOpenPopover = (component = <ContextMenuInPopover />) => {
+        cy.realMount(component);
+        cy.get('[data-test-subj="popoverToggle"]').click();
+        cy.wait(350); // EuiPopover's updateFocus() takes ~350ms to run
+      };
+
+      it('reclaims focus from the parent popover panel', () => {
+        mountAndOpenPopover();
         cy.focused().should('not.have.attr', 'class', 'euiPopover__panel');
         cy.focused().should('have.attr', 'class', 'euiContextMenuPanel');
       });
+
+      it('does not hijack focus from the EuiPopover if `initialFocus` is set', () => {
+        mountAndOpenPopover(
+          <ContextMenuInPopover initialFocus="#testInitialFocus">
+            <input id="testInitialFocus" />
+          </ContextMenuInPopover>
+        );
+        cy.focused().should('not.have.attr', 'class', 'euiContextMenuPanel');
+        cy.focused().should('have.attr', 'id', 'testInitialFocus');
+      });
+
+      it('restores focus to the toggling button on popover close', () => {
+        mountAndOpenPopover();
+        cy.realPress('Tab');
+        cy.realPress('Enter');
+        cy.focused().should('have.attr', 'data-test-subj', 'popoverToggle');
+      });
+
+      it('restores focus to the toggling button on popover escape key', () => {
+        mountAndOpenPopover();
+        cy.realPress('{esc}');
+        cy.focused().should('have.attr', 'data-test-subj', 'popoverToggle');
+      });
     });
   });
 

src/components/context_menu/context_menu_panel.tsx

@@ -84,6 +84,7 @@ interface State {
   focusedItemIndex?: number;
   currentHeight?: number;
   height?: number;
+  waitingForInitialPopover: boolean;
 }
 
 export class EuiContextMenuPanel extends Component<Props, State> {
@@ -109,6 +110,7 @@ export class EuiContextMenuPanel extends Component<Props, State> {
           ? props.initialFocusedItemIndex + 1 // Account for panel title back button
           : props.initialFocusedItemIndex,
       currentHeight: undefined,
+      waitingForInitialPopover: false,
     };
   }
 
@@ -226,6 +228,12 @@ export class EuiContextMenuPanel extends Component<Props, State> {
         return;
       }
 
+      // Don't take focus yet if EuiContextMenu is in a popover
+      // and the popover is initially opening/transitioning in
+      if (this.initialPopoverParent && this.state.waitingForInitialPopover) {
+        return;
+      }
+
       // Setting focus while transitioning causes the animation to glitch, so we have to wait
       // until it's finished before we focus anything.
       if (this.props.transitionType) {
@@ -265,16 +273,10 @@ export class EuiContextMenuPanel extends Component<Props, State> {
     });
   }
 
-  // If EuiContextMenu is used within an EuiPopover, EuiPopover's own
-  // `updateFocus()` method hijacks EuiContextMenuPanel's `updateFocus()`
-  // 350ms after the popover finishes transitioning in. This workaround
-  // reclaims focus from parent EuiPopovers that do not set an `initialFocus`
-  reclaimPopoverFocus() {
-    // If the popover panel gains focus, switch it to the context menu panel instead
-    this.initialPopoverParent?.addEventListener('focus', () => {
-      this.updateFocus();
-    });
-  }
+  reclaimPopoverFocus = () => {
+    this.setState({ waitingForInitialPopover: false });
+    this.updateFocus();
+  };
 
   onTransitionComplete = () => {
     if (this.props.onTransitionComplete) {
@@ -287,12 +289,28 @@ export class EuiContextMenuPanel extends Component<Props, State> {
   }
 
   componentDidMount() {
-    this.updateFocus();
-    this.reclaimPopoverFocus();
+    // If EuiContextMenu is used within an EuiPopover, we need to wait for EuiPopover to:
+    // 1. Correctly set its `returnFocus` to the toggling button,
+    //    so focus is correctly restored to the popover toggle on close
+    // 2. Finish its own `updateFocus()` call 350ms after transitioning in,
+    //    so the panel can handle its own focus without focus fighting
+    if (this.initialPopoverParent) {
+      this.initialPopoverParent.addEventListener(
+        'focus',
+        this.reclaimPopoverFocus,
+        { once: true }
+      );
+    } else {
+      this.updateFocus();
+    }
     this._isMounted = true;
   }
 
   componentWillUnmount() {
+    this.initialPopoverParent?.removeEventListener(
+      'focus',
+      this.reclaimPopoverFocus
+    );
     this._isMounted = false;
   }
 
@@ -365,6 +383,12 @@ export class EuiContextMenuPanel extends Component<Props, State> {
       return true;
     }
 
+    if (
+      nextState.waitingForInitialPopover !== this.state.waitingForInitialPopover
+    ) {
+      return true;
+    }
+
     // **
     // this component should have either items or children,
     // if there are items we can determine via `watchedItemProps` if we should update
@@ -424,6 +448,7 @@ export class EuiContextMenuPanel extends Component<Props, State> {
     if (!hasPopoverParent) return;
 
     this.initialPopoverParent = popoverParent;
+    this.setState({ waitingForInitialPopover: true });
   }
 
   panelRef = (node: HTMLElement | null) => {