fix(queryEnhancements): fix dataSource.id reference to enable tests (#7272) (#7406)

(cherry picked from commit c8496f8)

Signed-off-by: Joshua Li <joshuali925@gmail.com>
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Kawika Avilla <kavilla414@gmail.com>
Co-authored-by: Yulong Ruan <ruanyl@amazon.com>

Author: 4 people

src/plugins/query_enhancements/public/query_assist/components/query_assist_bar.tsx

@@ -45,7 +45,7 @@ export const QueryAssistBar: React.FC<QueryAssistInputProps> = (props) => {
     const subscription = props.connectionsService
       .getSelectedConnection$()
       .subscribe((connection) => {
-        dataSourceIdRef.current = connection?.id;
+        dataSourceIdRef.current = connection?.dataSource.id;
       });
     return () => subscription.unsubscribe();
   }, [props.connectionsService]);

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

@@ -24,13 +24,10 @@ const httpMock = coreSetupMock.http;
 
 jest.mock('../components', () => ({
   QueryAssistBar: jest.fn(() => <div>QueryAssistBar</div>),
-}));
-
-jest.mock('../components/query_assist_banner', () => ({
   QueryAssistBanner: jest.fn(() => <div>QueryAssistBanner</div>),
 }));
 
-describe.skip('CreateExtension', () => {
+describe('CreateExtension', () => {
   afterEach(() => {
     jest.clearAllMocks();
   });

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

@@ -28,7 +28,7 @@ const getAvailableLanguages$ = (
   connectionsService.getSelectedConnection$().pipe(
     distinctUntilChanged(),
     switchMap(async (connection) => {
-      const dataSourceId = connection?.id;
+      const dataSourceId = connection?.dataSource.id;
       const cached = availableLanguagesByDataSource.get(dataSourceId);
       if (cached !== undefined) return cached;
       const languages = await http

src/plugins/query_enhancements/server/routes/query_assist/agents.test.ts

@@ -5,25 +5,29 @@
 
 import { ApiResponse } from '@opensearch-project/opensearch';
 import { ResponseError } from '@opensearch-project/opensearch/lib/errors';
+import { loggerMock } from '@osd/logging/target/mocks';
 import { RequestHandlerContext } from 'src/core/server';
-// // eslint-disable-next-line @osd/eslint/no-restricted-paths
-// import { CoreRouteHandlerContext } from 'src/core/server/core_route_handler_context';
 import { coreMock } from '../../../../../core/server/mocks';
-import { loggerMock } from '@osd/logging/target/mocks';
 import { getAgentIdByConfig, requestAgentByConfig } from './agents';
 
-describe.skip('Agents helper functions', () => {
-  // const coreContext = new CoreRouteHandlerContext(
-  //   coreMock.createInternalStart(),
-  //   httpServerMock.createOpenSearchDashboardsRequest()
-  // );
+describe('Agents helper functions', () => {
   const coreContext = coreMock.createRequestHandlerContext();
   const client = coreContext.opensearch.client.asCurrentUser;
   const mockedTransport = client.transport.request as jest.Mock;
   const context: RequestHandlerContext = {
     core: coreContext,
+    // @ts-ignore
     dataSource: jest.fn(),
-    query_assist: { dataSourceEnabled: false, logger: loggerMock.create() },
+    query_assist: {
+      dataSourceEnabled: false,
+      logger: loggerMock.create(),
+      configPromise: Promise.resolve({
+        enabled: true,
+        queryAssist: {
+          supportedLanguages: [{ language: 'PPL', agentConfig: 'testConfig' }],
+        },
+      }),
+    },
   };
 
   afterEach(() => {

src/plugins/query_enhancements/server/types.ts

@@ -51,7 +51,10 @@ export interface FacetRequest {
   };
 }
 
-declare module '../../../core/server' {
+// TODO declaring it in core changes the interface for every reference. we only
+// need this declaration in query_enhancements, but it doesn't seem possible
+// https://github.com/opensearch-project/OpenSearch-Dashboards/issues/4274
+/* declare module '../../../core/server' {
   interface RequestHandlerContext {
     query_assist: {
       logger: Logger;
@@ -64,4 +67,4 @@ declare module '../../../core/server' {
       dataSourceEnabled: boolean;
     };
   }
-}
+} */