test: update tests

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

public/hooks/__tests__/use_conversations.test.ts

@@ -7,17 +7,36 @@ import { useDeleteConversation, usePatchConversation } from '../use_conversation
 import { renderHook, act } from '@testing-library/react-hooks';
 import { useCore } from '../../contexts/core_context';
 import { HttpHandler } from '../../../../../src/core/public';
-
+import { DataSourceServiceMock } from '../../services/data_source_service.mock';
+import { httpServiceMock } from '../../../../../src/core/public/mocks';
 jest.mock('../../contexts/core_context');
-const useCoreMocked = useCore as jest.MockedFunction<typeof useCore>;
+
+const setup = () => {
+  const useCoreMocked = useCore as jest.MockedFunction<typeof useCore>;
+  const mockedDataSource = new DataSourceServiceMock();
+  const mockedHttp = httpServiceMock.createStartContract();
+  mockedHttp.delete = jest.fn(() => Promise.resolve());
+  mockedHttp.put = jest.fn(() => Promise.resolve());
+
+  useCoreMocked.mockReturnValue({
+    services: {
+      dataSource: mockedDataSource,
+      http: mockedHttp,
+    },
+  });
+
+  return useCoreMocked;
+};
 
 describe('useDeleteConversation', () => {
+  const useCoreMocked = setup();
   it('should call delete with path and signal', async () => {
     const { result } = renderHook(() => useDeleteConversation());
 
     await act(async () => {
       await result.current.deleteConversation('foo');
     });
+
     expect(useCoreMocked.mock.results[0].value.services.http.delete).toHaveBeenCalledWith(
       '/api/assistant/conversation/foo',
       expect.objectContaining({
@@ -96,6 +115,8 @@ describe('useDeleteConversation', () => {
 });
 
 describe('usePatchConversation', () => {
+  const useCoreMocked = setup();
+
   it('should call put with path, query and signal', async () => {
     const { result } = renderHook(() => usePatchConversation());
 

public/hooks/use_chat_actions.test.tsx

@@ -13,6 +13,7 @@ import { ConversationLoadService } from '../services/conversation_load_service';
 import * as chatStateHookExports from './use_chat_state';
 import { ASSISTANT_API } from '../../common/constants/llm';
 import { IMessage } from 'common/types/chat_saved_object_attributes';
+import { DataSourceServiceMock } from '../services/data_source_service.mock';
 
 jest.mock('../services/conversations_service', () => {
   return {
@@ -60,6 +61,7 @@ describe('useChatActions hook', () => {
   const setSelectedTabIdMock = jest.fn();
   const pplVisualizationRenderMock = jest.fn();
   const setInteractionIdMock = jest.fn();
+  const dataSourceServiceMock = new DataSourceServiceMock();
 
   const chatContextMock: chatContextHookExports.IChatContext = {
     selectedTabId: 'chat',
@@ -88,8 +90,9 @@ describe('useChatActions hook', () => {
     jest.spyOn(coreHookExports, 'useCore').mockReturnValue({
       services: {
         http: httpMock,
-        conversations: new ConversationsService(httpMock),
-        conversationLoad: new ConversationLoadService(httpMock),
+        conversations: new ConversationsService(httpMock, dataSourceServiceMock),
+        conversationLoad: new ConversationLoadService(httpMock, dataSourceServiceMock),
+        dataSource: dataSourceServiceMock,
       },
     });
 
@@ -125,6 +128,7 @@ describe('useChatActions hook', () => {
         messages: [SEND_MESSAGE_RESPONSE.messages[0]],
         input: INPUT_MESSAGE,
       }),
+      query: dataSourceServiceMock.getDataSourceQuery(),
     });
 
     // it should send dispatch `receive` action to remove the message without messageId
@@ -199,6 +203,7 @@ describe('useChatActions hook', () => {
         messages: [],
         input: { type: 'input', content: 'message that send as input', contentType: 'text' },
       }),
+      query: dataSourceServiceMock.getDataSourceQuery(),
     });
   });
 
@@ -261,6 +266,7 @@ describe('useChatActions hook', () => {
     expect(chatStateDispatchMock).toHaveBeenCalledWith({ type: 'abort' });
     expect(httpMock.post).toHaveBeenCalledWith(ASSISTANT_API.ABORT_AGENT_EXECUTION, {
       body: JSON.stringify({ conversationId: 'conversation_id_to_abort' }),
+      query: dataSourceServiceMock.getDataSourceQuery(),
     });
   });
 
@@ -288,6 +294,7 @@ describe('useChatActions hook', () => {
         conversationId: 'conversation_id_mock',
         interactionId: 'interaction_id_mock',
       }),
+      query: dataSourceServiceMock.getDataSourceQuery(),
     });
     expect(chatStateDispatchMock).toHaveBeenCalledWith(
       expect.objectContaining({ type: 'receive', payload: { messages: [], interactions: [] } })
@@ -323,6 +330,7 @@ describe('useChatActions hook', () => {
         conversationId: 'conversation_id_mock',
         interactionId: 'interaction_id_mock',
       }),
+      query: dataSourceServiceMock.getDataSourceQuery(),
     });
     expect(chatStateDispatchMock).not.toHaveBeenCalledWith(
       expect.objectContaining({ type: 'receive' })

public/hooks/use_feed_back.test.tsx

@@ -12,11 +12,12 @@ import { httpServiceMock } from '../../../../src/core/public/mocks';
 import * as chatContextHookExports from '../contexts/chat_context';
 import { Interaction, IOutput, IMessage } from '../../common/types/chat_saved_object_attributes';
 import { ASSISTANT_API } from '../../common/constants/llm';
+import { DataSourceServiceMock } from '../services/data_source_service.mock';
 
 describe('useFeedback hook', () => {
   const httpMock = httpServiceMock.createStartContract();
   const chatStateDispatchMock = jest.fn();
-
+  const dataSourceMock = new DataSourceServiceMock();
   const chatContextMock = {
     rootAgentId: 'root_agent_id_mock',
     selectedTabId: 'chat',
@@ -29,6 +30,7 @@ describe('useFeedback hook', () => {
     jest.spyOn(coreHookExports, 'useCore').mockReturnValue({
       services: {
         http: httpMock,
+        dataSource: dataSourceMock,
       },
     });
     jest.spyOn(chatContextHookExports, 'useChatContext').mockReturnValue(chatContextMock);
@@ -82,6 +84,7 @@ describe('useFeedback hook', () => {
         body: JSON.stringify({
           satisfaction: true,
         }),
+        query: dataSourceMock.getDataSourceQuery(),
       }
     );
     expect(result.current.feedbackResult).toBe(true);
@@ -116,6 +119,7 @@ describe('useFeedback hook', () => {
         body: JSON.stringify({
           satisfaction: true,
         }),
+        query: dataSourceMock.getDataSourceQuery(),
       }
     );
     expect(result.current.feedbackResult).toBe(undefined);

public/services/__tests__/conversation_load_service.test.ts

@@ -6,10 +6,11 @@
 import { HttpHandler } from '../../../../../src/core/public';
 import { httpServiceMock } from '../../../../../src/core/public/mocks';
 import { ConversationLoadService } from '../conversation_load_service';
+import { DataSourceServiceMock } from '../data_source_service.mock';
 
 const setup = () => {
   const http = httpServiceMock.createSetupContract();
-  const conversationLoad = new ConversationLoadService(http);
+  const conversationLoad = new ConversationLoadService(http, new DataSourceServiceMock());
 
   return {
     conversationLoad,

public/services/__tests__/conversations_service.test.ts

@@ -6,14 +6,17 @@
 import { HttpHandler } from '../../../../../src/core/public';
 import { httpServiceMock } from '../../../../../src/core/public/mocks';
 import { ConversationsService } from '../conversations_service';
+import { DataSourceServiceMock } from '../data_source_service.mock';
 
 const setup = () => {
   const http = httpServiceMock.createSetupContract();
-  const conversations = new ConversationsService(http);
+  const dataSourceServiceMock = new DataSourceServiceMock();
+  const conversations = new ConversationsService(http, dataSourceServiceMock);
 
   return {
     conversations,
     http,
+    dataSource: dataSourceServiceMock,
   };
 };
 
@@ -32,7 +35,7 @@ describe('ConversationsService', () => {
   });
 
   it('should update options property and call get with passed query', () => {
-    const { conversations, http } = setup();
+    const { conversations, http, dataSource } = setup();
 
     expect(conversations.options).toBeFalsy();
     conversations.load({
@@ -49,6 +52,7 @@ describe('ConversationsService', () => {
         query: {
           page: 1,
           perPage: 10,
+          dataSourceId: dataSource.getDataSourceQuery()?.dataSourceId,
         },
         signal: expect.anything(),
       })
@@ -69,7 +73,7 @@ describe('ConversationsService', () => {
   });
 
   it('should call get with same query again after reload called', async () => {
-    const { conversations, http } = setup();
+    const { conversations, http, dataSource } = setup();
 
     await conversations.load({
       page: 1,
@@ -85,6 +89,7 @@ describe('ConversationsService', () => {
         query: {
           page: 1,
           perPage: 10,
+          dataSourceId: dataSource.getDataSourceQuery()?.dataSourceId,
         },
         signal: expect.anything(),
       })

public/services/data_source_service.mock.ts

@@ -0,0 +1,12 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+export class DataSourceServiceMock {
+  constructor() {}
+
+  getDataSourceQuery() {
+    return { dataSourceId: '' };
+  }
+}

public/services/data_source_service.ts

@@ -7,26 +7,10 @@ import { BehaviorSubject } from 'rxjs';
 import type { DataSourceManagementPluginSetup } from '../../../../src/plugins/data_source_management/public/plugin';
 
 export class DataSourceService {
-  dataSourceId$ = new BehaviorSubject<string | null>(null);
-  private dataSourceManagement: DataSourceManagementPluginSetup | undefined | null = null;
-
   constructor() {}
 
   getDataSourceQuery() {
-    if (!this.dataSourceManagement) {
-      return {};
-    }
     // TODO: use new handle logic to update
-    const dataSourceId = this.dataSourceManagement.dataSourceSelection
-      .getSelectionValue()
-      ?.values()
-      .next().value;
-    if (dataSourceId === null) {
-      throw new Error('No data source id');
-    }
-    if (dataSourceId === '') {
-      return {};
-    }
-    return { dataSourceId };
+    return { dataSourceId: '' };
   }
 }

public/tabs/history/__tests__/chat_history_page.test.tsx

@@ -14,6 +14,7 @@ import * as useChatStateExports from '../../../hooks/use_chat_state';
 import * as chatContextExports from '../../../contexts/chat_context';
 import * as coreContextExports from '../../../contexts/core_context';
 import { ConversationsService } from '../../../services/conversations_service';
+import { DataSourceServiceMock } from '../../../services/data_source_service.mock';
 
 import { ChatHistoryPage } from '../chat_history_page';
 
@@ -42,7 +43,7 @@ const setup = ({
     services: {
       ...coreMock.createStart(),
       http,
-      conversations: new ConversationsService(http),
+      conversations: new ConversationsService(http, new DataSourceServiceMock()),
       conversationLoad: {},
     },
   };

public/tabs/history/__tests__/chat_history_search_list.test.tsx

@@ -12,6 +12,7 @@ import * as chatContextExports from '../../../contexts/chat_context';
 import * as coreContextExports from '../../../contexts/core_context';
 
 import { ChatHistorySearchList, ChatHistorySearchListProps } from '../chat_history_search_list';
+import { DataSourceServiceMock } from '../../../services/data_source_service.mock';
 
 const setup = ({
   loading = false,
@@ -26,8 +27,9 @@ const setup = ({
     conversationId: '1',
     setTitle: jest.fn(),
   };
+  const dataSourceServiceMock = new DataSourceServiceMock();
   const useCoreMock = {
-    services: coreMock.createStart(),
+    services: { ...coreMock.createStart(), dataSource: dataSourceServiceMock },
   };
   useCoreMock.services.http.put.mockImplementation(() => Promise.resolve());
   useCoreMock.services.http.delete.mockImplementation(() => Promise.resolve());

public/tabs/history/__tests__/delete_conversation_confirm_modal.test.tsx

@@ -15,10 +15,13 @@ import {
   DeleteConversationConfirmModalProps,
 } from '../delete_conversation_confirm_modal';
 import { HttpHandler } from '../../../../../../src/core/public';
+import { DataSourceServiceMock } from '../../../services/data_source_service.mock';
 
 const setup = ({ onClose, conversationId }: DeleteConversationConfirmModalProps) => {
+  const dataSourceServiceMock = new DataSourceServiceMock();
+
   const useCoreMock = {
-    services: coreMock.createStart(),
+    services: { ...coreMock.createStart(), dataSource: dataSourceServiceMock },
   };
   jest.spyOn(coreContextExports, 'useCore').mockReturnValue(useCoreMock);