update mock context

Signed-off-by: tygao <tygao@amazon.com>
raintygao · May 28, 2024 · d1f9ce8 · d1f9ce8
1 parent 6c622df
commit d1f9ce8
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 27 deletions.
diff --git a/public/chat_flyout.tsx b/public/chat_flyout.tsx
@@ -26,10 +26,12 @@ export const ChatFlyout = (props: ChatFlyoutProps) => {
   const chatContext = useChatContext();
   const chatHistoryPageLoadedRef = useRef(false);
   const core = useCore();
+  const defaultDataSourceSelection$ = new BehaviorSubject(new Map());
+
   // TODO: use DataSourceService to replace
   const selectedDS = useObservable(
     core.services.setupDeps?.dataSourceManagement?.dataSourceSelection?.getSelection$() ??
-      new BehaviorSubject(new Map())
+      defaultDataSourceSelection$
   );
   const currentDS = useMemo(() => {
     return selectedDS?.values().next().value;

diff --git a/public/contexts/__mocks__/core_context.tsx b/public/contexts/__mocks__/core_context.tsx
@@ -24,6 +24,11 @@ export const useCore = jest.fn(() => {
         load: jest.fn(),
       },
       conversationLoad: {},
+      dataSource: {
+        getDataSourceQuery() {
+          return { dataSourceId: '' };
+        },
+      },
     },
   };
   useCoreMock.services.http.delete.mockReturnValue(Promise.resolve());

diff --git a/public/hooks/__tests__/use_conversations.test.ts b/public/hooks/__tests__/use_conversations.test.ts
@@ -7,36 +7,17 @@ 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 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;
-};
+jest.mock('../../contexts/core_context');
+const useCoreMocked = useCore as jest.MockedFunction<typeof useCore>;
 
 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({
@@ -115,8 +96,6 @@ describe('useDeleteConversation', () => {
 });
 
 describe('usePatchConversation', () => {
-  const useCoreMocked = setup();
-
   it('should call put with path, query and signal', async () => {
     const { result } = renderHook(() => usePatchConversation());
 

diff --git a/public/services/data_source_service.ts b/public/services/data_source_service.ts
@@ -3,9 +3,6 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { BehaviorSubject } from 'rxjs';
-import type { DataSourceManagementPluginSetup } from '../../../../src/plugins/data_source_management/public/plugin';
-
 export class DataSourceService {
   constructor() {}