Merge pull request #464 from Shopify/tf-inner-focus

[TrapFocus] Fix TrapFocus stealing focus from contained nodes

src/components/Focus/Focus.tsx

@@ -1,5 +1,6 @@
 import * as React from 'react';
 import * as ReactDOM from 'react-dom';
+import isEqual from 'lodash/isEqual';
 import {focusFirstFocusableNode} from '@shopify/javascript-utilities/focus';
 import {write} from '@shopify/javascript-utilities/fastdom';
 
@@ -10,6 +11,20 @@ export interface Props {
 
 export default class Focus extends React.PureComponent<Props, never> {
   componentDidMount() {
+    this.handleSelfFocus();
+  }
+
+  componentDidUpdate({children: prevChildren, ...restPrevProps}: Props) {
+    const {children, ...restProps} = this.props;
+
+    if (isEqual(restProps, restPrevProps)) {
+      return;
+    }
+
+    this.handleSelfFocus();
+  }
+
+  handleSelfFocus() {
     if (this.props.disabled) {
       return;
     }

src/components/TrapFocus/TrapFocus.tsx

@@ -14,14 +14,36 @@ export interface Props {
   children?: React.ReactNode;
 }
 
-export default class TrapFocus extends React.PureComponent<Props, never> {
-  private focusTrapWrapper: HTMLElement | null = null;
+export interface State {
+  shouldFocusSelf: boolean | undefined;
+}
+
+export default class TrapFocus extends React.PureComponent<Props, State> {
+  state = {
+    shouldFocusSelf: undefined,
+  };
+
+  private focusTrapWrapper: HTMLElement;
+
+  componentDidMount() {
+    this.setState(this.handleTrappingChange());
+  }
+
+  handleTrappingChange() {
+    const {trapping = true} = this.props;
+
+    if (this.focusTrapWrapper.contains(document.activeElement)) {
+      return {shouldFocusSelf: false};
+    }
+
+    return {shouldFocusSelf: trapping};
+  }
 
   render() {
-    const {children, trapping = true} = this.props;
+    const {children} = this.props;
 
     return (
-      <Focus disabled={!trapping}>
+      <Focus disabled={this.shouldDisable}>
         <div ref={this.setFocusTrapWrapper}>
           <EventListener event="focusout" handler={this.handleBlur} />
           {children}
@@ -30,6 +52,17 @@ export default class TrapFocus extends React.PureComponent<Props, never> {
     );
   }
 
+  private get shouldDisable() {
+    const {trapping = true} = this.props;
+    const {shouldFocusSelf} = this.state;
+
+    if (shouldFocusSelf === undefined) {
+      return true;
+    }
+
+    return shouldFocusSelf ? !trapping : !shouldFocusSelf;
+  }
+
   @autobind
   private setFocusTrapWrapper(node: HTMLDivElement) {
     this.focusTrapWrapper = node;

src/components/TrapFocus/tests/TrapFocus.test.tsx

@@ -1,9 +1,22 @@
 import * as React from 'react';
 import {mountWithAppProvider} from 'test-utilities';
-import {EventListener, Focus, TextContainer} from 'components';
+import {EventListener, Focus, TextContainer, TextField} from 'components';
+import {noop} from '@shopify/javascript-utilities/other';
 import TrapFocus from '../TrapFocus';
 
 describe('<TrapFocus />', () => {
+  let requestAnimationFrameSpy: jest.SpyInstance;
+
+  beforeEach(() => {
+    requestAnimationFrameSpy = jest.spyOn(window, 'requestAnimationFrame');
+    requestAnimationFrameSpy.mockImplementation((cb) => cb());
+  });
+
+  afterEach(() => {
+    (document.activeElement as HTMLElement).blur();
+    requestAnimationFrameSpy.mockRestore();
+  });
+
   it('mounts', () => {
     const trapFocus = mountWithAppProvider(
       <TrapFocus>
@@ -30,7 +43,6 @@ describe('<TrapFocus />', () => {
         <div />
       </TrapFocus>,
     ).find(Focus);
-
     expect(focus.prop('disabled')).toBe(false);
   });
 
@@ -53,4 +65,27 @@ describe('<TrapFocus />', () => {
 
     expect(focus.prop('disabled')).toBe(false);
   });
+
+  it('keeps focus on nodes contained inside trap focus during mount', () => {
+    const trapFocus = mountWithAppProvider(
+      <TrapFocus>
+        <TextField label="" value="" onChange={noop} autoFocus />
+      </TrapFocus>,
+    );
+    const input = trapFocus.find('input').getDOMNode();
+    expect(document.activeElement).toBe(input);
+  });
+
+  it('focuses the first focused node when nodes contained in trap focus are not in focus', () => {
+    const trapFocus = mountWithAppProvider(
+      <TrapFocus>
+        <a href="/">
+          <TextField label="" value="" onChange={noop} />
+        </a>
+      </TrapFocus>,
+    );
+    const focusedElement = trapFocus.find('a').getDOMNode();
+
+    expect(document.activeElement).toBe(focusedElement);
+  });
 });