From 1035ae2a4546b3272c5d5a2e6bac3b68b9094522 Mon Sep 17 00:00:00 2001
From: Kawika Avilla <kavilla414@gmail.com>
Date: Mon, 29 Jul 2024 05:00:04 +0000
Subject: [PATCH] fix mock for extension

Signed-off-by: Kawika Avilla <kavilla414@gmail.com>
---
 .../data/public/query/query_service.ts        |  4 +-
 .../utils/create_extension.test.tsx           | 47 ++++++++++++++-----
 2 files changed, 37 insertions(+), 14 deletions(-)

diff --git a/src/plugins/data/public/query/query_service.ts b/src/plugins/data/public/query/query_service.ts
index 346300dd1483..5c9c8b309367 100644
--- 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
index f79f646be314..c3fb47aef887 100644
--- 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', {