fix mock for extension

Signed-off-by: Kawika Avilla <[email protected]>
opensearch-project · Jul 29, 2024 · 1035ae2 · 1035ae2
1 parent b4e198a
commit 1035ae2
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 14 deletions.
diff --git a/src/plugins/data/public/query/query_service.ts b/src/plugins/data/public/query/query_service.ts
@@ -37,7 +37,7 @@ import { TimefilterService, TimefilterSetup } from './timefilter';
 import { createSavedQueryService } from './saved_query/saved_query_service';
 import { createQueryStateObservable } from './state_sync/create_global_query_observable';
 import { QueryStringManager, QueryStringContract } from './query_string';
-import { DataSetManager } from './dataset_manager';
+import { DataSetContract, DataSetManager } from './dataset_manager';
 import { buildOpenSearchQuery, getOpenSearchQueryConfig, IndexPatternsService } from '../../common';
 import { getUiSettings } from '../services';
 import { IndexPattern } from '..';
@@ -63,7 +63,7 @@ export class QueryService {
   filterManager!: FilterManager;
   timefilter!: TimefilterSetup;
   queryStringManager!: QueryStringContract;
-  dataSetManager!: DataSetManager;
+  dataSetManager!: DataSetContract;
 
   state$!: ReturnType<typeof createQueryStateObservable>;
 

diff --git a/src/plugins/query_enhancements/public/query_assist/utils/create_extension.test.tsx b/src/plugins/query_enhancements/public/query_assist/utils/create_extension.test.tsx
@@ -15,6 +15,24 @@ import { DataSetContract } from '../../../../data/public/query';
 import { ConfigSchema } from '../../../common/config';
 import { createQueryAssistExtension } from './create_extension';
 
+const mockSimpleDataSet = {
+  id: 'mock-data-set-id',
+  title: 'mock-title',
+  dataSourceRef: {
+    id: 'mock-data-source-id',
+  },
+} as SimpleDataSet;
+
+const mockDataSetManager: jest.Mocked<DataSetContract> = {
+  getUpdates$: jest.fn().mockReturnValue(of(mockSimpleDataSet)),
+  getDataSet: jest.fn().mockReturnValue(mockSimpleDataSet),
+  setDataSet: jest.fn(),
+  getDefaultDataSet: jest.fn().mockReturnValue(mockSimpleDataSet),
+  fetchDefaultDataSet: jest.fn().mockResolvedValue(mockSimpleDataSet),
+  init: jest.fn(),
+  initWithIndexPattern: jest.fn(),
+};
+
 const coreSetupMock = coreMock.createSetup({
   pluginStartDeps: {
     data: {
@@ -23,18 +41,13 @@ const coreSetupMock = coreMock.createSetup({
   },
 });
 const httpMock = coreSetupMock.http;
-const dataMock = dataPluginMock.createSetupContract();
-const dataSetMock = (dataMock.query.dataSetManager as unknown) as jest.Mocked<DataSetContract>;
-
-const mockSimpleDataSet = {
-  id: 'mock-data-set-id',
-  title: 'mock-title',
-  dataSourceRef: {
-    id: 'mock-data-source-id',
+const dataMock = {
+  ...dataPluginMock.createSetupContract(),
+  query: {
+    ...dataPluginMock.createSetupContract().query,
+    dataSetManager: mockDataSetManager,
   },
-} as SimpleDataSet;
-
-dataSetMock.getUpdates$.mockReturnValue(of(mockSimpleDataSet));
+};
 
 jest.mock('../components', () => ({
   QueryAssistBar: jest.fn(() => <div>QueryAssistBar</div>),
@@ -58,7 +71,17 @@ describe('CreateExtension', () => {
 
   it('should be enabled if at least one language is configured', async () => {
     httpMock.get.mockResolvedValueOnce({ configuredLanguages: ['PPL'] });
-    const extension = createQueryAssistExtension(httpMock, dataMock, config);
+    const extension = createQueryAssistExtension(
+      httpMock,
+      {
+        ...dataMock,
+        query: {
+          ...dataMock.query,
+          dataSetManager: mockDataSetManager,
+        },
+      },
+      config
+    );
     const isEnabled = await firstValueFrom(extension.isEnabled$(dependencies));
     expect(isEnabled).toBeTruthy();
     expect(httpMock.get).toBeCalledWith('/api/enhancements/assist/languages', {