enable query assist tests

Signed-off-by: Joshua Li <[email protected]>

src/plugins/data/public/query/query_string/dataset_service/lib/utils.test.ts

@@ -0,0 +1,82 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { LanguageServiceContract } from '../..';
+import { IQueryStart } from '../../..';
+import { DataStructure, DATA_STRUCTURE_META_TYPES } from '../../../../../common';
+import { dataPluginMock } from '../../../../mocks';
+import { setQueryService } from '../../../../services';
+import { injectMetaToDataStructures } from './utils';
+
+const mockDataStructures: DataStructure[] = [
+  {
+    id: 'fe25e2a0-6566-11ef-bb0e-0b6b1035facb',
+    title: 'mock-index-pattern-title',
+    type: 'INDEX_PATTERN',
+    parent: {
+      id: '8f26d980-63f5-11ef-b231-09f3ad4fb0e0',
+      title: 'mock-data-source-title',
+      type: 'OpenSearch',
+    },
+    meta: { type: DATA_STRUCTURE_META_TYPES.CUSTOM },
+  },
+];
+
+const dataMock = dataPluginMock.createSetupContract();
+const languageServiceMock = dataMock.query.queryString.getLanguageService() as jest.Mocked<
+  LanguageServiceContract
+>;
+setQueryService({ queryString: dataMock.query.queryString } as IQueryStart);
+
+languageServiceMock.getQueryEditorExtensionMap.mockReturnValue({
+  'mock-extension-1': {
+    id: 'mock-extension-1',
+    order: 1,
+    isEnabled$: jest.fn(),
+    getDataStructureMeta: (dataSourceId) =>
+      Promise.resolve({
+        type: DATA_STRUCTURE_META_TYPES.FEATURE,
+        icon: { type: 'icon1' },
+      }),
+  },
+  'mock-extension-2': {
+    id: 'mock-extension-2',
+    order: 2,
+    isEnabled$: jest.fn(),
+    getDataStructureMeta: (dataSourceId) =>
+      Promise.resolve({
+        type: DATA_STRUCTURE_META_TYPES.FEATURE,
+        icon: { type: 'icon2' },
+        tooltip: 'mock-extension-2',
+      }),
+  },
+});
+
+describe('Utils injectMetaToDataStructures', () => {
+  it('should inject meta', async () => {
+    const dataStructures = await injectMetaToDataStructures(mockDataStructures);
+    expect(dataStructures[0].meta).toMatchInlineSnapshot(`
+      Object {
+        "icon": Object {
+          "type": "icon1",
+        },
+        "tooltip": "mock-extension-2",
+        "type": "CUSTOM",
+      }
+    `);
+  });
+
+  it('does not change meta if not available', async () => {
+    languageServiceMock.getQueryEditorExtensionMap.mockReturnValue({
+      'mock-extension-3': {
+        id: 'mock-extension-3',
+        order: 3,
+        isEnabled$: jest.fn(),
+      },
+    });
+    const dataStructures = await injectMetaToDataStructures(mockDataStructures);
+    expect(dataStructures[0].meta).toBe(mockDataStructures[0].meta);
+  });
+});

src/plugins/query_enhancements/public/query_assist/utils/create_extension.test.tsx

@@ -6,11 +6,12 @@
 import { firstValueFrom } from '@osd/std';
 import { act, render, screen } from '@testing-library/react';
 import React from 'react';
+import { of } from 'rxjs';
 import { coreMock } from '../../../../../core/public/mocks';
-import { QueryEditorExtensionDependencies } from '../../../../data/public';
+import { QueryEditorExtensionDependencies, QueryStringContract } from '../../../../data/public';
 import { dataPluginMock } from '../../../../data/public/mocks';
 import { ConfigSchema } from '../../../common/config';
-import { createQueryAssistExtension } from './create_extension';
+import { clearCache, createQueryAssistExtension } from './create_extension';
 
 const coreSetupMock = coreMock.createSetup({
   pluginStartDeps: {
@@ -21,14 +22,32 @@ const coreSetupMock = coreMock.createSetup({
 });
 const httpMock = coreSetupMock.http;
 const dataMock = dataPluginMock.createSetupContract();
+const queryStringMock = dataMock.query.queryString as jest.Mocked<QueryStringContract>;
+
+const mockQueryWithIndexPattern = {
+  query: '',
+  language: 'kuery',
+  dataset: {
+    id: 'mock-index-pattern-id',
+    title: 'mock-index',
+    type: 'INDEX_PATTERN',
+    dataSource: {
+      id: 'mock-data-source-id',
+      title: 'test-mds',
+      type: 'OpenSearch',
+    },
+  },
+};
+
+queryStringMock.getQuery.mockReturnValue(mockQueryWithIndexPattern);
+queryStringMock.getUpdates$.mockReturnValue(of(mockQueryWithIndexPattern));
 
 jest.mock('../components', () => ({
   QueryAssistBar: jest.fn(() => <div>QueryAssistBar</div>),
   QueryAssistBanner: jest.fn(() => <div>QueryAssistBanner</div>),
 }));
 
-// TODO: https://github.com/opensearch-project/OpenSearch-Dashboards/issues/7860
-describe.skip('CreateExtension', () => {
+describe('CreateExtension', () => {
   const dependencies: QueryEditorExtensionDependencies = {
     language: 'PPL',
     onSelectLanguage: jest.fn(),
@@ -37,6 +56,7 @@ describe.skip('CreateExtension', () => {
   };
   afterEach(() => {
     jest.clearAllMocks();
+    clearCache();
   });
 
   const config: ConfigSchema['queryAssist'] = {
@@ -53,7 +73,7 @@ describe.skip('CreateExtension', () => {
     });
   });
 
-  it('should be disabled for unsupported language', async () => {
+  it('should be disabled when there is an error', async () => {
     httpMock.get.mockRejectedValueOnce(new Error('network failure'));
     const extension = createQueryAssistExtension(httpMock, dataMock, config);
     const isEnabled = await firstValueFrom(extension.isEnabled$(dependencies));
@@ -63,6 +83,35 @@ describe.skip('CreateExtension', () => {
     });
   });
 
+  it('creates data structure meta', async () => {
+    httpMock.get.mockResolvedValueOnce({ configuredLanguages: ['PPL'] });
+    const extension = createQueryAssistExtension(httpMock, dataMock, config);
+    const meta = await extension.getDataStructureMeta?.('mock-data-source-id2');
+    expect(meta).toMatchInlineSnapshot(`
+      Object {
+        "icon": Object {
+          "type": "test-file-stub",
+        },
+        "tooltip": "Query assist is available",
+        "type": "FEATURE",
+      }
+    `);
+    expect(httpMock.get).toBeCalledWith('/api/enhancements/assist/languages', {
+      query: { dataSourceId: 'mock-data-source-id2' },
+    });
+  });
+
+  it('does not send multiple requests for the same data source', async () => {
+    httpMock.get.mockResolvedValueOnce({ configuredLanguages: ['PPL'] });
+    const extension = createQueryAssistExtension(httpMock, dataMock, config);
+    const metas = await Promise.all(
+      Array.from({ length: 10 }, () => extension.getDataStructureMeta?.('mock-data-source-id2'))
+    );
+    metas.push(await extension.getDataStructureMeta?.('mock-data-source-id2'));
+    metas.forEach((meta) => expect(meta?.type).toBe('FEATURE'));
+    expect(httpMock.get).toBeCalledTimes(1);
+  });
+
   it('should render the component if language is supported', async () => {
     httpMock.get.mockResolvedValueOnce({ configuredLanguages: ['PPL'] });
     const extension = createQueryAssistExtension(httpMock, dataMock, config);

src/plugins/query_enhancements/public/query_assist/utils/create_extension.tsx

@@ -17,35 +17,41 @@ import { ConfigSchema } from '../../../common/config';
 import { QueryAssistBanner, QueryAssistBar } from '../components';
 import assistantMark from '../../assets/query_assist_mark.svg';
 
-/**
- * @returns list of query assist supported languages for the given data source.
- */
-const getAvailableLanguagesForDataSource = (() => {
+const [getAvailableLanguagesForDataSource, clearCache] = (() => {
   const availableLanguagesByDataSource: Map<string | undefined, string[]> = new Map();
   const pendingRequests: Map<string | undefined, Promise<string[]>> = new Map();
 
-  return async (http: HttpSetup, dataSourceId: string | undefined) => {
-    const cached = availableLanguagesByDataSource.get(dataSourceId);
-    if (cached !== undefined) return cached;
-
-    const pendingRequest = pendingRequests.get(dataSourceId);
-    if (pendingRequest !== undefined) return pendingRequest;
-
-    const languagesPromise = http
-      .get<{ configuredLanguages: string[] }>(API.QUERY_ASSIST.LANGUAGES, {
-        query: { dataSourceId },
-      })
-      .then((response) => response.configuredLanguages)
-      .catch(() => [])
-      .finally(() => pendingRequests.delete(dataSourceId));
-    pendingRequests.set(dataSourceId, languagesPromise);
-
-    const languages = await languagesPromise;
-    availableLanguagesByDataSource.set(dataSourceId, languages);
-    return languages;
-  };
+  return [
+    async (http: HttpSetup, dataSourceId: string | undefined) => {
+      const cached = availableLanguagesByDataSource.get(dataSourceId);
+      if (cached !== undefined) return cached;
+
+      const pendingRequest = pendingRequests.get(dataSourceId);
+      if (pendingRequest !== undefined) return pendingRequest;
+
+      const languagesPromise = http
+        .get<{ configuredLanguages: string[] }>(API.QUERY_ASSIST.LANGUAGES, {
+          query: { dataSourceId },
+        })
+        .then((response) => response.configuredLanguages)
+        .catch(() => [])
+        .finally(() => pendingRequests.delete(dataSourceId));
+      pendingRequests.set(dataSourceId, languagesPromise);
+
+      const languages = await languagesPromise;
+      availableLanguagesByDataSource.set(dataSourceId, languages);
+      return languages;
+    },
+    () => {
+      availableLanguagesByDataSource.clear();
+      pendingRequests.clear();
+    },
+  ];
 })();
 
+// visible for testing
+export { clearCache };
+
 /**
  * @returns observable list of query assist agent configured languages in the
  * selected data source.