Connect: Improve focus management in search bar

web/packages/teleterm/src/ui/Search/SearchBar.test.tsx

@@ -15,19 +15,24 @@
  */
 
 import React from 'react';
-import { render, screen } from 'design/utils/testing';
+import { render, screen, waitFor } from 'design/utils/testing';
 import { makeSuccessAttempt } from 'shared/hooks/useAsync';
 
 import { MockAppContext } from 'teleterm/ui/fixtures/mocks';
 import { MockAppContextProvider } from 'teleterm/ui/fixtures/MockAppContextProvider';
 import { ResourceSearchError } from 'teleterm/ui/services/resources';
+import ModalsHost from 'teleterm/ui/ModalsHost';
 
 import * as pickers from './pickers/pickers';
 import * as useActionAttempts from './pickers/useActionAttempts';
 import * as SearchContext from './SearchContext';
 
 import { SearchBarConnected } from './SearchBar';
 
+beforeEach(() => {
+  jest.restoreAllMocks();
+});
+
 it('does not display empty results copy after selecting two filters', () => {
   const appContext = new MockAppContext();
   appContext.workspacesService.setState(draft => {
@@ -202,6 +207,63 @@ it('notifies about resource search errors and allows to display details', () =>
   expect(mockedSearchContext.lockOpen).toHaveBeenCalled();
 });
 
+it('maintains focus on the search input after closing a resource search error modal', async () => {
+  const appContext = new MockAppContext();
+  appContext.workspacesService.setState(draft => {
+    draft.rootClusterUri = '/clusters/foo';
+  });
+
+  const resourceSearchError = new ResourceSearchError(
+    '/clusters/foo',
+    'server',
+    new Error('whoops')
+  );
+
+  const mockActionAttempts = {
+    filterActionsAttempt: makeSuccessAttempt([]),
+    resourceActionsAttempt: makeSuccessAttempt([]),
+    resourceSearchAttempt: makeSuccessAttempt({
+      results: [],
+      errors: [resourceSearchError],
+      search: '',
+    }),
+  };
+  jest
+    .spyOn(useActionAttempts, 'useActionAttempts')
+    .mockImplementation(() => mockActionAttempts);
+
+  render(
+    <MockAppContextProvider appContext={appContext}>
+      <SearchBarConnected />
+      <ModalsHost />
+    </MockAppContextProvider>
+  );
+
+  screen.getByRole('searchbox').focus();
+  expect(screen.getByRole('menu')).toHaveTextContent(
+    'Some of the search results are incomplete.'
+  );
+  screen.getByText('Show details').click();
+
+  const modal = screen.getByTestId('Modal');
+  expect(modal).toHaveTextContent('Resource search errors');
+  expect(modal).toHaveTextContent('whoops');
+
+  // Lose focus on the search input.
+  screen.getByText('Close').focus();
+  screen.getByText('Close').click();
+
+  // Need to await this since some state updates in SearchContext are done after the modal closes.
+  // Otherwise we'd get a warning about missing `act`.
+  await waitFor(() => {
+    expect(modal).not.toBeInTheDocument();
+  });
+
+  expect(screen.getByRole('searchbox')).toHaveFocus();
+  // Verify that the search bar wasn't closed.
+  expect(screen.getByRole('menu')).toBeInTheDocument();
+});
+
 const getMockedSearchContext = () => ({
   inputValue: 'foo',
   filters: [],

web/packages/teleterm/src/ui/Search/SearchContext.test.tsx

@@ -15,6 +15,8 @@
  */
 
 import React from 'react';
+import '@testing-library/jest-dom';
+import { render, screen } from '@testing-library/react';
 import { renderHook, act } from '@testing-library/react-hooks';
 
 import { SearchContextProvider, useSearchContext } from './SearchContext';
@@ -40,11 +42,16 @@ describe('lockOpen', () => {
       finishAction: rejectFailureAction,
     },
   ])('$name', async ({ action, finishAction }) => {
+    const inputFocus = jest.fn();
     const { result } = renderHook(() => useSearchContext(), {
       wrapper: ({ children }) => (
         <SearchContextProvider>{children}</SearchContextProvider>
       ),
     });
+    result.current.inputRef.current = {
+      focus: inputFocus,
+    } as unknown as HTMLInputElement;
+
     act(() => {
       result.current.open();
     });
@@ -74,5 +81,44 @@ describe('lockOpen', () => {
       result.current.close();
     });
     expect(result.current.isOpen).toBe(false);
+
+    // Called the first time during `open`, then again after `lockOpen` finishes.
+    expect(inputFocus).toHaveBeenCalledTimes(2);
+  });
+});
+
+describe('open', () => {
+  it('manages the focus properly when called with no arguments', () => {
+    const SearchInput = () => {
+      const { inputRef, open, close } = useSearchContext();
+
+      return (
+        <>
+          <input data-testid="search-input" ref={inputRef} />
+          <button data-testid="open" onClick={() => open()} />
+          <button data-testid="close" onClick={() => close()} />
+        </>
+      );
+    };
+
+    render(
+      <>
+        <input data-testid="other-input" />
+
+        <SearchContextProvider>
+          <SearchInput />
+        </SearchContextProvider>
+      </>
+    );
+
+    const otherInput = screen.getByTestId('other-input');
+    const searchInput = screen.getByTestId('search-input');
+    otherInput.focus();
+
+    screen.getByTestId('open').click();
+    expect(searchInput).toHaveFocus();
+
+    screen.getByTestId('close').click();
+    expect(otherInput).toHaveFocus();
   });
 });

web/packages/teleterm/src/ui/Search/SearchContext.tsx

@@ -95,14 +95,24 @@ export const SearchContextProvider: FC = props => {
 
   function open(fromElement?: Element): void {
     if (isOpen) {
+      // Even if the search bar is already open, we want to focus on the input anyway. The search
+      // input might lose focus due to user interaction while the search bar stays open. Focusing
+      // here again makes it possible to use the shortcut to grant the focus to the input again.
+      inputRef.current?.focus();
       return;
     }
+
+    // In case `open` was called without `fromElement` (e.g. when using the keyboard shortcut), we
+    // must read `document.activeElement` before we focus the input.
     previouslyActive.current = fromElement || document.activeElement;
+    inputRef.current?.focus();
     setIsOpen(true);
   }
 
   /**
-   * lockOpen forces the search bar to stay open for the duration of the action.
+   * lockOpen forces the search bar to stay open for the duration of the action. It also restores
+   * focus on the search input after the action is done.
+   *
    * This is useful in situations where want the search bar to not close when the user interacts
    * with modals shown from the search bar.
    */
@@ -112,6 +122,12 @@ export const SearchContextProvider: FC = props => {
     try {
       await action;
     } finally {
+      // By the time the action passes, the user might have caused the focus to be lost on the
+      // search input, so let's bring it back.
+      //
+      // focus needs to be executed before the state update, otherwise the search bar will close for
+      // some reason.
+      inputRef.current?.focus();
       setIsLockedOpen(false);
     }
   }

web/packages/teleterm/src/ui/Search/useSearch.test.tsx

@@ -33,6 +33,10 @@ import { rankResults, useFilterSearch, useResourceSearch } from './useSearch';
 
 import type * as tsh from 'teleterm/services/tshd/types';
 
+beforeEach(() => {
+  jest.restoreAllMocks();
+});
+
 describe('rankResults', () => {
   it('uses the displayed resource name as the tie breaker if the scores are equal', () => {
     const server = makeResourceResult({