test: add test for saving to notebook

Signed-off-by: tygao <[email protected]>

public/components/chat_window_header_title.tsx

@@ -25,6 +25,7 @@ export const ChatWindowHeaderTitle = React.memo(() => {
   const core = useCore();
   const [isPopoverOpen, setPopoverOpen] = useState(false);
   const [isRenameModalOpen, setRenameModalOpen] = useState(false);
+  const [isSaveNotebookModalOpen, setSaveNotebookModalOpen] = useState(false);
   const { chatState } = useChatState();
   const { saveChat } = useSaveChat();
 
@@ -50,6 +51,10 @@ export const ChatWindowHeaderTitle = React.memo(() => {
     [chatContext, core.services.sessions]
   );
 
+  const handleSaveNotebookModalClose = () => {
+    setSaveNotebookModalOpen(false);
+  };
+
   const button = (
     <EuiFlexGroup
       style={{ maxWidth: '300px', padding: '0 8px' }}
@@ -100,10 +105,8 @@ export const ChatWindowHeaderTitle = React.memo(() => {
     <EuiContextMenuItem
       key="save-as-notebook"
       onClick={() => {
-        const modal = core.overlays.openModal(
-          <NotebookNameModal onClose={() => modal.close()} saveChat={saveChat} />
-        );
         closePopover();
+        setSaveNotebookModalOpen(true);
       }}
       // User only can save conversation when he send a message at least.
       disabled={chatState.messages.every((item) => item.type !== 'input')}
@@ -131,6 +134,9 @@ export const ChatWindowHeaderTitle = React.memo(() => {
           defaultTitle={chatContext.title!}
         />
       )}
+      {isSaveNotebookModalOpen && (
+        <NotebookNameModal onClose={handleSaveNotebookModalClose} saveChat={saveChat} />
+      )}
     </>
   );
 });

public/components/notebook/__tests__/notebook_name_modal.test.tsx

@@ -0,0 +1,115 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import React from 'react';
+import { act, fireEvent, render, waitFor } from '@testing-library/react';
+import { I18nProvider } from '@osd/i18n/react';
+
+import { coreMock } from '../../../../../../src/core/public/mocks';
+import * as coreContextExports from '../../../contexts/core_context';
+
+import { NotebookNameModal, NotebookNameModalProps } from '../notebook_name_modal';
+import { AssistantServices } from '../../../contexts/core_context';
+import { OpenSearchDashboardsReactContextValue } from '../../../../../../src/plugins/opensearch_dashboards_react/public';
+
+const setup = ({ onClose, saveChat }: NotebookNameModalProps) => {
+  const useCoreMock = {
+    services: coreMock.createStart(),
+  };
+  jest.spyOn(coreContextExports, 'useCore').mockReturnValue(
+    // In test env, only mock needed core service and assert.
+    (useCoreMock as unknown) as OpenSearchDashboardsReactContextValue<AssistantServices>
+  );
+
+  const renderResult = render(
+    <I18nProvider>
+      <NotebookNameModal onClose={onClose} saveChat={saveChat} />
+    </I18nProvider>
+  );
+
+  return {
+    useCoreMock,
+    renderResult,
+  };
+};
+
+describe('<NotebookNameModal />', () => {
+  it('should call onClose after cancel button click', async () => {
+    const onCloseMock = jest.fn();
+    const saveChatMock = jest.fn();
+    const { renderResult, useCoreMock } = setup({
+      onClose: onCloseMock,
+      saveChat: saveChatMock,
+    });
+
+    expect(onCloseMock).not.toHaveBeenCalled();
+
+    act(() => {
+      fireEvent.click(renderResult.getByTestId('confirmNotebookCancelButton'));
+    });
+
+    await waitFor(() => {
+      expect(onCloseMock).toHaveBeenCalled();
+    });
+  });
+
+  it('should show success toast and call onClose after saving chat succeed', async () => {
+    const onCloseMock = jest.fn();
+    const saveChatMock = jest.fn();
+    const { renderResult, useCoreMock } = setup({
+      onClose: onCloseMock,
+      saveChat: saveChatMock,
+    });
+
+    act(() => {
+      fireEvent.change(renderResult.getByLabelText('Notebook name input'), {
+        target: {
+          value: 'notebook-name',
+        },
+      });
+    });
+
+    expect(onCloseMock).not.toHaveBeenCalled();
+
+    act(() => {
+      fireEvent.click(renderResult.getByTestId('confirmNotebookConfirmButton'));
+    });
+
+    await waitFor(() => {
+      expect(useCoreMock.services.notifications.toasts.addSuccess).toHaveBeenCalled();
+      expect(onCloseMock).toHaveBeenCalled();
+    });
+  });
+
+  it('should show error toasts and call onClose after failed save chat', async () => {
+    const onCloseMock = jest.fn();
+    const saveChatMock = jest.fn(() => {
+      throw new Error();
+    });
+    const { renderResult, useCoreMock } = setup({
+      onClose: onCloseMock,
+      saveChat: saveChatMock,
+    });
+
+    act(() => {
+      fireEvent.change(renderResult.getByLabelText('Notebook name input'), {
+        target: {
+          value: 'notebook-name',
+        },
+      });
+    });
+
+    expect(onCloseMock).not.toHaveBeenCalled();
+
+    act(() => {
+      fireEvent.click(renderResult.getByTestId('confirmNotebookConfirmButton'));
+    });
+
+    await waitFor(() => {
+      expect(useCoreMock.services.notifications.toasts.addDanger).toHaveBeenCalled();
+      expect(onCloseMock).toHaveBeenCalled();
+    });
+  });
+});

public/components/notebook/notebook_name_modal.tsx

@@ -6,7 +6,7 @@
 import {
   EuiButton,
   EuiButtonEmpty,
-  EuiForm,
+  EuiLink,
   EuiFormRow,
   EuiModal,
   EuiModalBody,
@@ -16,22 +16,54 @@ import {
   EuiModalHeaderTitle,
 } from '@elastic/eui';
 import React, { useState, useCallback } from 'react';
+import { useCore } from '../../contexts/core_context';
+import { toMountPoint } from '../../../../../src/plugins/opensearch_dashboards_react/public';
 
-interface Props {
+export interface NotebookNameModalProps {
   onClose: () => void;
   // SaveChat hook depends on context. Runtime modal component can't get context, so saveChat needs to be passed in.
   saveChat: (name: string) => void;
 }
 
-export const NotebookNameModal = ({ onClose, saveChat }: Props) => {
+export const NotebookNameModal = ({ onClose, saveChat }: NotebookNameModalProps) => {
+  const {
+    services: {
+      notifications: { toasts },
+    },
+  } = useCore();
   const [name, setName] = useState('');
   const [loading, setLoading] = useState(false);
 
   const onSubmit = useCallback(async () => {
     setLoading(true);
-    await saveChat(name);
+    try {
+      const notebookId = await saveChat(name);
+      const notebookLink = `./observability-notebooks#/${notebookId}?view=view_both`;
+
+      toasts.addSuccess({
+        text: toMountPoint(
+          <>
+            <p>
+              This conversation was saved as{' '}
+              <EuiLink href={notebookLink} target="_blank">
+                {name}
+              </EuiLink>
+              .
+            </p>
+          </>
+        ),
+      });
+    } catch (error) {
+      if (error.message === 'Not Found') {
+        toasts.addDanger(
+          'This feature depends on the observability plugin, please install it before use.'
+        );
+      } else {
+        toasts.addDanger('Failed to save to notebook');
+      }
+    }
     onClose();
-  }, [name, saveChat, onclose]);
+  }, [name, saveChat, onclose, toasts.addSuccess, toasts.addDanger]);
 
   return (
     <>
@@ -42,18 +74,25 @@ export const NotebookNameModal = ({ onClose, saveChat }: Props) => {
 
         <EuiModalBody>
           <EuiFormRow label="Please enter a name for your notebook.">
-            <EuiFieldText value={name} onChange={(e) => setName(e.target.value)} />
+            <EuiFieldText
+              value={name}
+              onChange={(e) => setName(e.target.value)}
+              aria-label="Notebook name input"
+            />
           </EuiFormRow>
         </EuiModalBody>
 
         <EuiModalFooter>
-          <EuiButtonEmpty onClick={onClose}>Cancel</EuiButtonEmpty>
+          <EuiButtonEmpty onClick={onClose} data-test-subj="confirmNotebookCancelButton">
+            Cancel
+          </EuiButtonEmpty>
           <EuiButton
             type="submit"
             fill
             isLoading={loading}
             disabled={name.length < 1}
             onClick={onSubmit}
+            data-test-subj="confirmNotebookConfirmButton"
           >
             Confirm name
           </EuiButton>

public/hooks/use_save_chat.tsx

@@ -5,13 +5,10 @@
 
 import React from 'react';
 import { useCallback } from 'react';
-import { EuiLink } from '@elastic/eui';
 import { NOTEBOOK_API } from '../../common/constants/llm';
 import { useCore } from '../contexts/core_context';
 import { useChatState } from './use_chat_state';
 import { convertMessagesToParagraphs, Paragraphs } from '../utils';
-import { getCoreStart } from '../plugin';
-import { toMountPoint } from '../../../../src/plugins/opensearch_dashboards_react/public';
 import { useChatContext } from '../contexts/chat_context';
 
 interface SetParagraphResponse {
@@ -62,40 +59,13 @@ export const useSaveChat = () => {
 
   const saveChat = useCallback(
     async (name: string) => {
-      try {
-        const id = await createNotebook(name);
-        const paragraphs = convertMessagesToParagraphs(
-          chatState.messages,
-          chatContext.currentAccount.username
-        );
-        await setParagraphs(id, paragraphs);
-        const notebookLink = `./observability-notebooks#/${id}?view=view_both`;
-
-        getCoreStart().notifications.toasts.addSuccess({
-          text: toMountPoint(
-            <>
-              <p>
-                This conversation was saved as{' '}
-                <EuiLink href={notebookLink} target="_blank">
-                  {name}
-                </EuiLink>
-                .
-              </p>
-            </>
-          ),
-        });
-      } catch (error) {
-        if (error.message === 'Not Found') {
-          getCoreStart().notifications.toasts.addError(error, {
-            title:
-              'This feature depends on the observability plugin, please install it before use.',
-          });
-        } else {
-          getCoreStart().notifications.toasts.addError(error, {
-            title: 'Failed to save to notebook',
-          });
-        }
-      }
+      const id = await createNotebook(name);
+      const paragraphs = convertMessagesToParagraphs(
+        chatState.messages,
+        chatContext.currentAccount.username
+      );
+      await setParagraphs(id, paragraphs);
+      return id;
     },
     [chatState, createNotebook, setParagraphs, chatContext]
   );