add unit tests for HeaderChatButton component

public/chat_flyout.tsx

@@ -21,7 +21,7 @@ interface ChatFlyoutProps {
   toggleFlyoutFullScreen: () => void;
 }
 
-export const ChatFlyout: React.FC<ChatFlyoutProps> = (props) => {
+export const ChatFlyout = (props: ChatFlyoutProps) => {
   const chatContext = useChatContext();
   const chatHistoryPageLoadedRef = useRef(false);
 

public/chat_header_button.test.tsx

@@ -0,0 +1,162 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import React from 'react';
+import { act, render, fireEvent, screen } from '@testing-library/react';
+
+import { HeaderChatButton } from './chat_header_button';
+import { applicationServiceMock } from '../../../src/core/public/mocks';
+import { AssistantActions } from './types';
+import { BehaviorSubject } from 'rxjs';
+
+let mockSend: jest.Mock;
+let mockLoadChat: jest.Mock;
+
+jest.mock('./hooks/use_chat_actions', () => {
+  mockSend = jest.fn();
+  mockLoadChat = jest.fn();
+  return {
+    useChatActions: jest.fn().mockReturnValue({
+      send: mockSend,
+      loadChat: mockLoadChat,
+      openChatUI: jest.fn(),
+      executeAction: jest.fn(),
+      abortAction: jest.fn(),
+      regenerate: jest.fn(),
+    }),
+  };
+});
+
+jest.mock('./chat_flyout', () => {
+  return {
+    ChatFlyout: ({
+      toggleFlyoutFullScreen,
+      flyoutFullScreen,
+    }: {
+      toggleFlyoutFullScreen: () => void;
+      flyoutFullScreen: boolean;
+    }) => (
+      <div aria-label="chat flyout mock">
+        <button onClick={toggleFlyoutFullScreen}>toggle chat flyout fullscreen</button>
+        <p>{flyoutFullScreen ? 'fullscreen mode' : 'dock-right mode'}</p>
+      </div>
+    ),
+  };
+});
+
+describe('<HeaderChatButton />', () => {
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it('should open chat flyout and send the initial message', () => {
+    const applicationStart = {
+      ...applicationServiceMock.createStartContract(),
+      currentAppId$: new BehaviorSubject(''),
+    };
+    render(
+      <HeaderChatButton
+        application={applicationStart}
+        userHasAccess={true}
+        contentRenderers={{}}
+        actionExecutors={{}}
+        assistantActions={{} as AssistantActions}
+        currentAccount={{ username: 'test_user', tenant: 'test_tenant' }}
+      />
+    );
+
+    act(() => applicationStart.currentAppId$.next('mock_app_id'));
+
+    screen.getByLabelText('chat input').focus();
+    fireEvent.change(screen.getByLabelText('chat input'), {
+      target: { value: 'what indices are in my cluster?' },
+    });
+    expect(screen.getByLabelText('chat input')).toHaveFocus();
+
+    fireEvent.keyPress(screen.getByLabelText('chat input'), {
+      key: 'Enter',
+      code: 'Enter',
+      charCode: 13,
+    });
+
+    // start a new chat
+    expect(mockLoadChat).toHaveBeenCalled();
+    // send chat message
+    expect(mockSend).toHaveBeenCalledWith({
+      type: 'input',
+      contentType: 'text',
+      content: 'what indices are in my cluster?',
+      context: { appId: 'mock_app_id' },
+    });
+    // chat flyout displayed
+    expect(screen.queryByLabelText('chat flyout mock')).toBeInTheDocument();
+    // the input value is cleared after pressing enter
+    expect(screen.getByLabelText('chat input')).toHaveValue('');
+    expect(screen.getByLabelText('chat input')).not.toHaveFocus();
+  });
+
+  it('should toggle chat flyout size', () => {
+    render(
+      <HeaderChatButton
+        application={applicationServiceMock.createStartContract()}
+        userHasAccess={true}
+        contentRenderers={{}}
+        actionExecutors={{}}
+        assistantActions={{} as AssistantActions}
+        currentAccount={{ username: 'test_user', tenant: 'test_tenant' }}
+      />
+    );
+    fireEvent.click(screen.getByLabelText('toggle chat flyout icon'));
+    expect(screen.queryByText('dock-right mode')).toBeInTheDocument();
+
+    fireEvent.click(screen.getByText('toggle chat flyout fullscreen'));
+    expect(screen.queryByText('fullscreen mode')).toBeInTheDocument();
+  });
+
+  it('should focus in chat input when click and press Escape should blur', () => {
+    render(
+      <HeaderChatButton
+        application={applicationServiceMock.createStartContract()}
+        userHasAccess={true}
+        contentRenderers={{}}
+        actionExecutors={{}}
+        assistantActions={{} as AssistantActions}
+        currentAccount={{ username: 'test_user', tenant: 'test_tenant' }}
+      />
+    );
+    screen.getByLabelText('chat input').focus();
+    expect(screen.getByLabelText('chat input')).toHaveFocus();
+    expect(screen.getByTitle('press enter to chat')).toBeInTheDocument();
+
+    fireEvent.keyUp(screen.getByLabelText('chat input'), {
+      key: 'Escape',
+      code: 'Escape',
+      charCode: 27,
+    });
+    expect(screen.getByLabelText('chat input')).not.toHaveFocus();
+    expect(screen.getByTitle('press ⌘ + / to start typing')).toBeInTheDocument();
+  });
+
+  it('should focus on chat input when pressing global shortcut', () => {
+    render(
+      <HeaderChatButton
+        application={applicationServiceMock.createStartContract()}
+        userHasAccess={true}
+        contentRenderers={{}}
+        actionExecutors={{}}
+        assistantActions={{} as AssistantActions}
+        currentAccount={{ username: 'test_user', tenant: 'test_tenant' }}
+      />
+    );
+    expect(screen.getByLabelText('chat input')).not.toHaveFocus();
+    fireEvent.keyDown(document.body, {
+      key: '/',
+      code: 'NumpadDivide',
+      charCode: 111,
+      metaKey: true,
+    });
+    expect(screen.getByLabelText('chat input')).toHaveFocus();
+  });
+});

public/chat_header_button.tsx

@@ -28,7 +28,7 @@ interface HeaderChatButtonProps {
 
 let flyoutLoaded = false;
 
-export const HeaderChatButton: React.FC<HeaderChatButtonProps> = (props) => {
+export const HeaderChatButton = (props: HeaderChatButtonProps) => {
   const [appId, setAppId] = useState<string>();
   const [sessionId, setSessionId] = useState<string>();
   const [title, setTitle] = useState<string>();
@@ -144,6 +144,7 @@ export const HeaderChatButton: React.FC<HeaderChatButtonProps> = (props) => {
     <>
       <div className={classNames('llm-chat-header-icon-wrapper')}>
         <EuiFieldText
+          aria-label="chat input"
           inputRef={inputRef}
           compressed
           value={query}
@@ -154,16 +155,29 @@ export const HeaderChatButton: React.FC<HeaderChatButtonProps> = (props) => {
           onKeyPress={onKeyPress}
           onKeyUp={onKeyUp}
           prepend={
-            <EuiIcon type={chatIcon} size="l" onClick={() => setFlyoutVisible(!flyoutVisible)} />
+            <EuiIcon
+              aria-label="toggle chat flyout icon"
+              type={chatIcon}
+              size="l"
+              onClick={() => setFlyoutVisible(!flyoutVisible)}
+            />
           }
           append={
             <span className="llm-chat-header-shortcut">
               {inputFocus ? (
-                <EuiBadge className="llm-chat-header-shortcut-enter" color="hollow">
+                <EuiBadge
+                  title="press enter to chat"
+                  className="llm-chat-header-shortcut-enter"
+                  color="hollow"
+                >
                   ⏎
                 </EuiBadge>
               ) : (
-                <EuiBadge className="llm-chat-header-shortcut-cmd" color="hollow">
+                <EuiBadge
+                  title="press ⌘ + / to start typing"
+                  className="llm-chat-header-shortcut-cmd"
+                  color="hollow"
+                >
                   ⌘ + /
                 </EuiBadge>
               )}

public/contexts/set_context.tsx

@@ -12,7 +12,7 @@ interface SetContextProps {
 }
 
 // TODO needs a better solution to expose hook
-export const SetContext: React.FC<SetContextProps> = (props) => {
+export const SetContext = (props: SetContextProps) => {
   Object.assign(props.assistantActions, useChatActions());
   return null;
 };

public/hooks/use_chat_actions.test.tsx

@@ -13,21 +13,18 @@ import { SessionLoadService } from '../services/session_load_service';
 import * as chatStateHookExports from './use_chat_state';
 import { ASSISTANT_API } from '../../common/constants/llm';
 
-const mockSessionsLoad = jest.fn();
-const mockSessionLoad = jest.fn().mockReturnValue({ messages: [], interactions: [] });
-
 jest.mock('../services/sessions_service', () => {
   return {
     SessionsService: jest.fn().mockImplementation(() => {
-      return { reload: mockSessionsLoad };
+      return { reload: jest.fn() };
     }),
   };
 });
 
 jest.mock('../services/session_load_service', () => {
   return {
     SessionLoadService: jest.fn().mockImplementation(() => {
-      return { load: mockSessionLoad };
+      return { load: jest.fn().mockReturnValue({ messages: [], interactions: [] }) };
     }),
   };
 });

public/hooks/use_chat_state.tsx

@@ -79,7 +79,7 @@ const chatStateReducer: React.Reducer<ChatState, ChatStateAction> = (state, acti
     }
   });
 
-export const ChatStateProvider: React.FC = (props) => {
+export const ChatStateProvider = (props: { children?: React.ReactNode }) => {
   const [chatState, chatStateDispatch] = useReducer(chatStateReducer, initialState);
   const contextValue: IChatStateContext = useMemo(() => ({ chatState, chatStateDispatch }), [
     chatState,

public/tabs/chat_window_header.tsx

@@ -15,7 +15,7 @@ export interface ChatWindowHeaderProps {
   toggleFlyoutFullScreen: () => void;
 }
 
-export const ChatWindowHeader: React.FC<ChatWindowHeaderProps> = React.memo((props) => {
+export const ChatWindowHeader = React.memo((props: ChatWindowHeaderProps) => {
   const chatContext = useChatContext();
 
   const dockBottom = () => (