From 641408baf4c64ae7da16daa6bb201d14bc00d503 Mon Sep 17 00:00:00 2001
From: Joshua Li <joshuali925@gmail.com>
Date: Mon, 29 Jul 2024 22:55:28 +0000
Subject: [PATCH] simplify tests

Signed-off-by: Joshua Li <joshuali925@gmail.com>
---
 .../dataset_manager/dataset_manager.mock.ts   |  2 +
 src/plugins/data/public/query/mocks.ts        |  4 +-
 .../create_dataset_navigator.tsx              |  2 +-
 .../utils/create_extension.test.tsx           | 48 +++++--------------
 4 files changed, 18 insertions(+), 38 deletions(-)

diff --git a/src/plugins/data/public/query/dataset_manager/dataset_manager.mock.ts b/src/plugins/data/public/query/dataset_manager/dataset_manager.mock.ts
index 2f1f5144274..93dcdfb22b8 100644
--- a/src/plugins/data/public/query/dataset_manager/dataset_manager.mock.ts
+++ b/src/plugins/data/public/query/dataset_manager/dataset_manager.mock.ts
@@ -12,6 +12,8 @@ const createSetupContractMock = () => {
     setDataSet: jest.fn(),
     getUpdates$: jest.fn(),
     getDefaultDataSet: jest.fn(),
+    fetchDefaultDataSet: jest.fn(),
+    initWithIndexPattern: jest.fn(),
   };
   return dataSetManagerMock;
 };
diff --git a/src/plugins/data/public/query/mocks.ts b/src/plugins/data/public/query/mocks.ts
index 200ef46a517..87bd2edb49b 100644
--- a/src/plugins/data/public/query/mocks.ts
+++ b/src/plugins/data/public/query/mocks.ts
@@ -42,7 +42,7 @@ const createSetupContractMock = () => {
     filterManager: createFilterManagerMock(),
     timefilter: timefilterServiceMock.createSetupContract(),
     queryString: queryStringManagerMock.createSetupContract(),
-    dataSet: dataSetManagerMock.createSetupContract(),
+    dataSetManager: dataSetManagerMock.createSetupContract(),
     state$: new Observable(),
   };
 
@@ -57,7 +57,7 @@ const createStartContractMock = () => {
     savedQueries: jest.fn() as any,
     state$: new Observable(),
     timefilter: timefilterServiceMock.createStartContract(),
-    dataSet: dataSetManagerMock.createStartContract(),
+    dataSetManager: dataSetManagerMock.createStartContract(),
     getOpenSearchQuery: jest.fn(),
   };
 
diff --git a/src/plugins/data/public/ui/dataset_navigator/create_dataset_navigator.tsx b/src/plugins/data/public/ui/dataset_navigator/create_dataset_navigator.tsx
index 94674ab2bcf..b6a6274e13c 100644
--- a/src/plugins/data/public/ui/dataset_navigator/create_dataset_navigator.tsx
+++ b/src/plugins/data/public/ui/dataset_navigator/create_dataset_navigator.tsx
@@ -15,7 +15,7 @@ export function createDataSetNavigator(
   dataSetManager: DataSetContract
 ) {
   // Return a function that takes props, omitting the dependencies from the props type
-  return (props: Omit<DataSetNavigatorProps, 'savedObjectsClient' | 'http' | 'dataSet'>) => (
+  return (props: Omit<DataSetNavigatorProps, 'savedObjectsClient' | 'http' | 'dataSetManager'>) => (
     <DataSetNavigator
       {...props}
       savedObjectsClient={savedObjectsClient}
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 c3fb47aef88..cf99c79798b 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,24 +15,6 @@ 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: {
@@ -41,13 +23,19 @@ const coreSetupMock = coreMock.createSetup({
   },
 });
 const httpMock = coreSetupMock.http;
-const dataMock = {
-  ...dataPluginMock.createSetupContract(),
-  query: {
-    ...dataPluginMock.createSetupContract().query,
-    dataSetManager: mockDataSetManager,
+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',
   },
-};
+} as SimpleDataSet;
+
+dataSetMock.getDataSet.mockReturnValue(mockSimpleDataSet);
+dataSetMock.getUpdates$.mockReturnValue(of(mockSimpleDataSet));
 
 jest.mock('../components', () => ({
   QueryAssistBar: jest.fn(() => <div>QueryAssistBar</div>),
@@ -71,17 +59,7 @@ describe('CreateExtension', () => {
 
   it('should be enabled if at least one language is configured', async () => {
     httpMock.get.mockResolvedValueOnce({ configuredLanguages: ['PPL'] });
-    const extension = createQueryAssistExtension(
-      httpMock,
-      {
-        ...dataMock,
-        query: {
-          ...dataMock.query,
-          dataSetManager: mockDataSetManager,
-        },
-      },
-      config
-    );
+    const extension = createQueryAssistExtension(httpMock, dataMock, config);
     const isEnabled = await firstValueFrom(extension.isEnabled$(dependencies));
     expect(isEnabled).toBeTruthy();
     expect(httpMock.get).toBeCalledWith('/api/enhancements/assist/languages', {