[TrapFocus] Fixed focus theft

UNRELEASED.md

@@ -8,6 +8,7 @@
 
 ### Bug fixes
 
+- Fixed `TrapFocus` stealing focus from other `TrapFocus`'s ([#2681](https://github.com/Shopify/polaris-react/pull/2681))
 - Fixed focus state color on monochrome `Buttons` ([#2684](https://github.com/Shopify/polaris-react/pull/2684))
 
 ### Documentation

src/components/AppProvider/AppProvider.tsx

@@ -2,6 +2,7 @@ import React from 'react';
 import {ThemeConfig} from '../../utilities/theme';
 import {ThemeProvider} from '../ThemeProvider';
 import {MediaQueryProvider} from '../MediaQueryProvider';
+import {FocusManager} from '../FocusManager';
 import {I18n, I18nContext} from '../../utilities/i18n';
 import {
   ScrollLockManager,
@@ -110,7 +111,9 @@ export class AppProvider extends React.Component<AppProviderProps, State> {
                 <AppBridgeContext.Provider value={appBridge}>
                   <LinkContext.Provider value={link}>
                     <ThemeProvider theme={theme}>
-                      <MediaQueryProvider>{children}</MediaQueryProvider>
+                      <MediaQueryProvider>
+                        <FocusManager>{children}</FocusManager>
+                      </MediaQueryProvider>
                     </ThemeProvider>
                   </LinkContext.Provider>
                 </AppBridgeContext.Provider>

src/components/AppProvider/tests/AppProvider.test.tsx

@@ -6,6 +6,7 @@ import {mountWithApp} from 'test-utilities/react-testing';
 import {MediaQueryProvider} from 'components/MediaQueryProvider';
 import {LinkContext} from '../../../utilities/link';
 import {AppProvider} from '../AppProvider';
+import {FocusManager} from '../../FocusManager';
 
 describe('<AppProvider />', () => {
   beforeEach(() => {
@@ -42,4 +43,13 @@ describe('<AppProvider />', () => {
     );
     expect(appProvider).toContainReactComponent(MediaQueryProvider);
   });
+
+  it('renders a FocusManager', () => {
+    const appProvider = mountWithApp(
+      <AppProvider i18n={{}}>
+        <div>Child</div>
+      </AppProvider>,
+    );
+    expect(appProvider).toContainReactComponent(FocusManager);
+  });
 });

src/components/FocusManager/FocusManager.tsx

@@ -0,0 +1,42 @@
+import React, {useMemo, useState, useCallback, ContextType} from 'react';
+import {FocusManagerContext} from '../../utilities/focus-manager';
+
+interface Props {
+  children?: React.ReactNode;
+}
+
+type Context = NonNullable<ContextType<typeof FocusManagerContext>>;
+
+export function FocusManager({children}: Props) {
+  const [trapFocusList, setTrapFocusList] = useState<string[]>([]);
+
+  const add = useCallback<Context['add']>((id) => {
+    setTrapFocusList((list) => [...list, id]);
+  }, []);
+  const remove = useCallback<Context['remove']>((id) => {
+    let removed = true;
+    setTrapFocusList((list) => {
+      const clone = [...list];
+      const index = clone.indexOf(id);
+      if (index === -1) {
+        removed = false;
+      } else {
+        clone.splice(index, 1);
+      }
+      return clone;
+    });
+    return removed;
+  }, []);
+
+  const value = useMemo(() => ({trapFocusList, add, remove}), [
+    add,
+    trapFocusList,
+    remove,
+  ]);
+
+  return (
+    <FocusManagerContext.Provider value={value}>
+      {children}
+    </FocusManagerContext.Provider>
+  );
+}

src/components/FocusManager/index.ts

@@ -0,0 +1 @@
+export * from './FocusManager';

src/components/FocusManager/tests/FocusManager.test.tsx

@@ -0,0 +1,67 @@
+import React, {useContext, useEffect, useRef} from 'react';
+import {mountWithApp} from 'test-utilities';
+import {
+  useFocusManager,
+  FocusManagerContext,
+} from '../../../utilities/focus-manager';
+
+const Component = ({id}: {id: string}) =>
+  useFocusManager().canSafelyFocus ? <div id={id} /> : null;
+
+describe('FocusManager', () => {
+  it('allows the first component added to be safely focused', () => {
+    const id = 'one';
+    const component = mountWithApp(
+      <div>
+        <Component id={id} />
+        <Component id="two" />
+      </div>,
+    );
+    expect(component).toContainReactComponentTimes('div', 1, {id});
+  });
+
+  it('does not allow the second component added to be safely focused', () => {
+    const id = 'two';
+    const component = mountWithApp(
+      <div>
+        <Component id="one" />
+        <Component id={id} />
+      </div>,
+    );
+    expect(component).not.toContainReactComponent('div', {id});
+  });
+
+  describe('remove', () => {
+    it('returns false when the component was not removed', () => {
+      const Component = () => {
+        const wasRemoved = useRef(false);
+        const {remove} = useContext(FocusManagerContext)!;
+
+        useEffect(() => {
+          wasRemoved.current = remove('id');
+        }, [remove]);
+
+        return wasRemoved.current ? <div /> : null;
+      };
+      const component = mountWithApp(<Component />);
+      expect(component).not.toContainReactComponent('div');
+    });
+
+    it('returns true when the component was added', () => {
+      const id = 'id';
+      const Component = () => {
+        const wasRemoved = useRef(false);
+        const {add, remove} = useContext(FocusManagerContext)!;
+
+        useEffect(() => {
+          add(id);
+          wasRemoved.current = remove(id);
+        }, [add, remove]);
+
+        return wasRemoved.current ? <div /> : null;
+      };
+      const component = mountWithApp(<Component />);
+      expect(component).toContainReactComponentTimes('div', 1);
+    });
+  });
+});

src/components/PolarisTestProvider/PolarisTestProvider.tsx

@@ -1,4 +1,5 @@
 import React from 'react';
+import {FocusManager} from '../FocusManager';
 import {merge} from '../../utilities/merge';
 import {FrameContext} from '../../utilities/frame';
 import {
@@ -104,9 +105,11 @@ export function PolarisTestProvider({
                   <LinkContext.Provider value={link}>
                     <ThemeContext.Provider value={mergedTheme}>
                       <MediaQueryContext.Provider value={mergedMediaQuery}>
-                        <FrameContext.Provider value={mergedFrame}>
-                          {children}
-                        </FrameContext.Provider>
+                        <FocusManager>
+                          <FrameContext.Provider value={mergedFrame}>
+                            {children}
+                          </FrameContext.Provider>
+                        </FocusManager>
                       </MediaQueryContext.Provider>
                     </ThemeContext.Provider>
                   </LinkContext.Provider>

src/components/TrapFocus/TrapFocus.tsx

@@ -12,6 +12,7 @@ import {
   findLastKeyboardFocusableNode,
   focusLastKeyboardFocusableNode,
 } from '../../utilities/focus';
+import {useFocusManager} from '../../utilities/focus-manager';
 
 export interface TrapFocusProps {
   trapping?: boolean;
@@ -22,20 +23,21 @@ export function TrapFocus({trapping = true, children}: TrapFocusProps) {
   const [shouldFocusSelf, setFocusSelf] = useState<boolean | undefined>(
     undefined,
   );
-
+  const {canSafelyFocus} = useFocusManager();
   const focusTrapWrapper = useRef<HTMLDivElement>(null);
 
   useEffect(() => {
     setFocusSelf(
       !(
+        canSafelyFocus &&
         focusTrapWrapper.current &&
         focusTrapWrapper.current.contains(document.activeElement)
       ),
     );
-  }, []);
+  }, [canSafelyFocus]);
 
   const shouldDisableFirstElementFocus = () => {
-    if (shouldFocusSelf === undefined) {
+    if (shouldFocusSelf === undefined || !canSafelyFocus) {
       return true;
     }
 
@@ -56,6 +58,7 @@ export function TrapFocus({trapping = true, children}: TrapFocusProps) {
     }
 
     if (
+      canSafelyFocus &&
       focusTrapWrapper.current !== event.target &&
       !focusTrapWrapper.current.contains(event.target as Node)
     ) {

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

@@ -124,6 +124,22 @@ describe('<TrapFocus />', () => {
     expect(document.activeElement).toBe(focusedElement);
   });
 
+  it(`doesn't trade steal focus from another TrapFocus when multiple are rendered`, () => {
+    const id = 'input';
+    const trapFocus = mountWithApp(
+      <div>
+        <TrapFocus>
+          <input autoFocus id={id} />
+        </TrapFocus>
+        <TrapFocus>
+          <input autoFocus />
+        </TrapFocus>
+      </div>,
+    );
+
+    expect(trapFocus.find('input', {id})!.domNode).toBe(document.activeElement);
+  });
+
   describe('handleBlur', () => {
     const externalDomNode = mountWithAppProvider(<Button />)
       .find('button')

src/utilities/focus-manager/context.tsx

@@ -0,0 +1,11 @@
+import React from 'react';
+
+export interface FocusManagerContextType {
+  trapFocusList: string[];
+  add: (id: string) => void;
+  remove: (id: string) => boolean;
+}
+
+export const FocusManagerContext = React.createContext<
+  FocusManagerContextType | undefined
+>(undefined);

src/utilities/focus-manager/hooks.ts

@@ -0,0 +1,30 @@
+import {useContext, useMemo, useEffect} from 'react';
+import {useUniqueId} from '../unique-id';
+import {MissingAppProviderError} from '../errors';
+import {FocusManagerContext} from './context';
+
+export function useFocusManager() {
+  const focusManager = useContext(FocusManagerContext);
+  const id = useUniqueId();
+
+  if (!focusManager) {
+    throw new MissingAppProviderError('No FocusManager was provided.');
+  }
+
+  const {
+    trapFocusList,
+    add: addFocusItem,
+    remove: removeFocusItem,
+  } = focusManager;
+  const canSafelyFocus = trapFocusList[0] === id;
+  const value = useMemo(() => ({canSafelyFocus}), [canSafelyFocus]);
+
+  useEffect(() => {
+    addFocusItem(id);
+    return () => {
+      removeFocusItem(id);
+    };
+  }, [addFocusItem, id, removeFocusItem]);
+
+  return value;
+}

src/utilities/focus-manager/index.ts

@@ -0,0 +1,2 @@
+export * from './context';
+export * from './hooks';

src/utilities/focus-manager/tests/hooks.test.tsx

@@ -0,0 +1,42 @@
+import React from 'react';
+import {mount, mountWithApp} from 'test-utilities';
+import {useFocusManager} from '../hooks';
+import {
+  UniqueIdFactory,
+  UniqueIdFactoryContext,
+  globalIdGeneratorFactory,
+} from '../../unique-id';
+
+let consoleErrorSpy: jest.SpyInstance;
+
+const Component = () =>
+  typeof useFocusManager().canSafelyFocus === 'boolean' ? <div /> : null;
+
+describe('useFocusManager', () => {
+  beforeEach(() => {
+    consoleErrorSpy = jest.spyOn(console, 'error').mockImplementation(() => {});
+  });
+
+  afterEach(() => {
+    consoleErrorSpy.mockRestore();
+  });
+
+  it('returns context', () => {
+    const component = mountWithApp(<Component />);
+    expect(component).toContainReactComponent('div');
+  });
+
+  it('throws an error if context is not set', () => {
+    const attemptMount = () =>
+      mount(
+        <UniqueIdFactoryContext.Provider
+          value={new UniqueIdFactory(globalIdGeneratorFactory)}
+        >
+          <Component />
+        </UniqueIdFactoryContext.Provider>,
+      );
+    expect(attemptMount).toThrow(
+      'No FocusManager was provided. Your application must be wrapped in an <AppProvider> component. See https://polaris.shopify.com/components/structure/app-provider for implementation instructions.',
+    );
+  });
+});