From 6d89db1307ce2735c30268bd8128b9044f0198ab Mon Sep 17 00:00:00 2001
From: Andrew Musgrave <andrew.musgrave@shopify.com>
Date: Mon, 7 Oct 2019 17:27:05 -0400
Subject: [PATCH] Fixed Popover focusing incorrect element when closed

---
 src/components/Popover/Popover.tsx            |   5 +-
 src/components/Popover/tests/Popover.test.tsx |  41 +++++++
 src/utilities/focus.ts                        |  34 ++++++
 src/utilities/tests/focus.test.ts             | 106 ++++++++++++++++--
 4 files changed, 175 insertions(+), 11 deletions(-)

diff --git a/src/components/Popover/Popover.tsx b/src/components/Popover/Popover.tsx
index 2d06a55dcb4..1d11951bc28 100644
--- a/src/components/Popover/Popover.tsx
+++ b/src/components/Popover/Popover.tsx
@@ -4,6 +4,7 @@ import {
   focusFirstFocusableNode,
   findFirstFocusableNode,
 } from '@shopify/javascript-utilities/focus';
+import {focusNextFocusableNode} from '../../utilities/focus';
 
 import {PreferredPosition, PreferredAlignment} from '../PositionedOverlay';
 import {Portal} from '../Portal';
@@ -135,7 +136,9 @@ export class Popover extends React.PureComponent<PopoverProps, State> {
       source === CloseSource.FocusOut ||
       source === CloseSource.EscapeKeypress
     ) {
-      focusFirstFocusableNode(this.activatorContainer, false);
+      if (!focusNextFocusableNode(this.activatorContainer)) {
+        focusFirstFocusableNode(this.activatorContainer, false);
+      }
     }
   };
 
diff --git a/src/components/Popover/tests/Popover.test.tsx b/src/components/Popover/tests/Popover.test.tsx
index 980036fdf60..93c5ef46208 100644
--- a/src/components/Popover/tests/Popover.test.tsx
+++ b/src/components/Popover/tests/Popover.test.tsx
@@ -1,6 +1,8 @@
 import React from 'react';
 import {mountWithAppProvider, findByTestID} from 'test-utilities/legacy';
+import {mountWithApp} from 'test-utilities';
 import {Popover} from '../Popover';
+import {PopoverOverlay} from '../components';
 
 describe('<Popover />', () => {
   const spy = jest.fn();
@@ -182,4 +184,43 @@ describe('<Popover />', () => {
 
     expect(onCloseSpy).not.toHaveBeenCalled();
   });
+
+  it('focuses the next available element when the popover is closed', () => {
+    const id = 'focus-target';
+    function PopoverTest() {
+      return (
+        <React.Fragment>
+          <Popover active activator={<div />} onClose={noop} />
+          <button id={id} />
+        </React.Fragment>
+      );
+    }
+
+    const popover = mountWithApp(<PopoverTest />);
+
+    popover.find(PopoverOverlay)!.trigger('onClose', 1);
+    const focusTarget = popover.find('button', {id})!.domNode;
+
+    expect(document.activeElement).toBe(focusTarget);
+  });
+
+  it('focuses the activator when there is another focusable element is not available when the popover is closed', () => {
+    const id = 'activator';
+    function PopoverTest() {
+      return (
+        <React.Fragment>
+          <Popover active activator={<button id={id} />} onClose={noop} />
+        </React.Fragment>
+      );
+    }
+
+    const popover = mountWithApp(<PopoverTest />);
+
+    popover.find(PopoverOverlay)!.trigger('onClose', 1);
+    const focusTarget = popover.find('button', {id})!.domNode;
+
+    expect(document.activeElement).toBe(focusTarget);
+  });
 });
+
+function noop() {}
diff --git a/src/utilities/focus.ts b/src/utilities/focus.ts
index 1c4b9b2e0ca..cb5d370c899 100644
--- a/src/utilities/focus.ts
+++ b/src/utilities/focus.ts
@@ -1,5 +1,39 @@
+import {
+  findFirstFocusableNode,
+  FOCUSABLE_SELECTOR,
+} from '@shopify/javascript-utilities/focus';
+
 export function handleMouseUpByBlurring({
   currentTarget,
 }: React.MouseEvent<HTMLAnchorElement | HTMLButtonElement>) {
   currentTarget.blur();
 }
+
+export function nextFocusableNode(
+  node: HTMLElement,
+): HTMLElement | Element | null {
+  let nextElement = node.nextElementSibling;
+
+  while (nextElement) {
+    if (nextElement.matches(FOCUSABLE_SELECTOR)) return nextElement;
+
+    if (nextElement instanceof HTMLElement) {
+      const innerFocusableElement = findFirstFocusableNode(nextElement);
+      if (innerFocusableElement) return innerFocusableElement;
+    }
+
+    nextElement = nextElement.nextElementSibling;
+  }
+
+  return nextElement;
+}
+
+export function focusNextFocusableNode(node: HTMLElement) {
+  const nextFocusable = nextFocusableNode(node);
+  if (nextFocusable && nextFocusable instanceof HTMLElement) {
+    nextFocusable.focus();
+    return true;
+  }
+
+  return false;
+}
diff --git a/src/utilities/tests/focus.test.ts b/src/utilities/tests/focus.test.ts
index d71cbeedc69..7d1f6b12697 100644
--- a/src/utilities/tests/focus.test.ts
+++ b/src/utilities/tests/focus.test.ts
@@ -1,14 +1,100 @@
 import {MouseEvent} from 'react';
-import {handleMouseUpByBlurring} from '../focus';
-
-describe('focus', () => {
-  describe('handleMouseUpByBlurring()', () => {
-    it('calls blur on the currentTarget', () => {
-      const currentTarget = document.createElement('button');
-      jest.spyOn(currentTarget, 'blur');
-      const mouseEvent = {currentTarget};
-      handleMouseUpByBlurring(mouseEvent as MouseEvent<HTMLButtonElement>);
-      expect(currentTarget.blur).toHaveBeenCalled();
+import {
+  handleMouseUpByBlurring,
+  focusNextFocusableNode,
+  nextFocusableNode,
+} from '../focus';
+
+describe('handleMouseUpByBlurring()', () => {
+  it('calls blur on the currentTarget', () => {
+    const currentTarget = document.createElement('button');
+    jest.spyOn(currentTarget, 'blur');
+    const mouseEvent = {currentTarget};
+    handleMouseUpByBlurring(mouseEvent as MouseEvent<HTMLButtonElement>);
+    expect(currentTarget.blur).toHaveBeenCalled();
+  });
+});
+
+describe('nextFocusableNode', () => {
+  it('does not return the initial element as the focusable node', () => {
+    const {activator, otherNode} = domSetup();
+
+    expect(nextFocusableNode(activator)).toBe(otherNode);
+  });
+
+  it('returns null when a focusable element is not found', () => {
+    const {activator} = domSetup({
+      otherNodeTag: 'div',
     });
+
+    expect(nextFocusableNode(activator)).toBeNull();
+  });
+
+  it("returns the parent of an adjacent element when it's focusable", () => {
+    const {activator, otherNode} = domSetup();
+
+    expect(nextFocusableNode(activator)).toBe(otherNode);
+  });
+
+  it('searches adjacent elements for focusable children', () => {
+    const {activator, otherNodeNested} = domSetup({
+      otherNodeTag: 'div',
+      nested: true,
+    });
+
+    expect(nextFocusableNode(activator)).toBe(otherNodeNested);
   });
 });
+
+describe('focusNextFocusableNode', () => {
+  it('returns true when the node was focused', () => {
+    const {activator} = domSetup();
+
+    expect(focusNextFocusableNode(activator)).toBe(true);
+  });
+
+  it('returns false when the node was not focused', () => {
+    const {activator} = domSetup({otherNodeTag: 'div'});
+
+    expect(focusNextFocusableNode(activator)).toBe(false);
+  });
+
+  it('focused the node', () => {
+    const {activator, otherNode} = domSetup();
+
+    focusNextFocusableNode(activator);
+
+    expect(document.activeElement).toBe(otherNode);
+  });
+});
+
+function domSetup(
+  options: {
+    wrapperTag?: string;
+    activatorTag?: string;
+    otherNodeTag?: string;
+    otherNodeNestedTag?: string;
+    nested?: boolean;
+  } = {},
+) {
+  const {
+    wrapperTag = 'div',
+    activatorTag = 'button',
+    otherNodeTag = 'button',
+    otherNodeNestedTag = 'button',
+    nested,
+  } = options;
+  const wrapper = document.createElement(wrapperTag);
+  const activator = document.createElement(activatorTag);
+  const otherNode = document.createElement(otherNodeTag);
+  let otherNodeNested = null;
+
+  if (nested) {
+    otherNodeNested = document.createElement(otherNodeNestedTag);
+    otherNode.appendChild(otherNodeNested);
+  }
+
+  wrapper.append(activator, otherNode);
+
+  return {wrapper, activator, otherNode, otherNodeNested};
+}