Address comments

Signed-off-by: Liyun Xiu <[email protected]>

src/plugins/data/public/ui/query_editor/query_editor.tsx

@@ -14,10 +14,9 @@ import {
   EuiText,
   PopoverAnchorPosition,
 } from '@elastic/eui';
-import { BehaviorSubject } from 'rxjs';
 import classNames from 'classnames';
 import { isEqual } from 'lodash';
-import React, { Component, createRef, RefObject } from 'react';
+import React, { Component, createRef, RefObject, useMemo } from 'react';
 import { monaco } from '@osd/monaco';
 import {
   IDataPluginServices,
@@ -338,16 +337,22 @@ export default class QueryEditorUI extends Component<Props, State> {
   };
 
   private renderExtensionSearchBarButtion = () => {
-    const supported = 'query-assist';
     if (!this.extensionMap || Object.keys(this.extensionMap).length === 0) return null;
-    if (!(supported in this.extensionMap) || !this.extensionMap[supported].getSearchBarButton)
-      return null;
-    return this.extensionMap[supported].getSearchBarButton({
-      language: this.props.query.language,
-      onSelectLanguage: this.onSelectLanguage,
-      isCollapsed: this.state.isCollapsed,
-      setIsCollapsed: this.setIsCollapsed,
-    });
+    const sortedConfigs = Object.values(this.extensionMap).sort((a, b) => a.order - b.order);
+    return (
+      <>
+        {sortedConfigs.map((config) => {
+          return config.getSearchBarButton
+            ? config.getSearchBarButton({
+                language: this.props.query.language,
+                onSelectLanguage: this.onSelectLanguage,
+                isCollapsed: this.state.isCollapsed,
+                setIsCollapsed: this.setIsCollapsed,
+              })
+            : null;
+        })}
+      </>
+    );
   };
 
   public render() {

src/plugins/query_enhancements/common/query_assist/index.ts

@@ -3,4 +3,4 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-export { QueryAssistParameters, QueryAssistResponse } from './types';
+export { QueryAssistParameters, QueryAssistResponse, QueryAssistContextType } from './types';

src/plugins/query_enhancements/common/query_assist/types.ts

@@ -17,3 +17,9 @@ export interface QueryAssistParameters {
   // for MDS
   dataSourceId?: string;
 }
+
+export enum QueryAssistContextType {
+  QUESTION,
+  QUERY,
+  DATA,
+}

src/plugins/query_enhancements/public/query_assist/_index.scss

@@ -8,6 +8,17 @@
     margin-top: $euiSizeXS;
   }
 
+  &.queryAssist__summary_banner {
+    /* stylelint-disable  @osd/stylelint/no_restricted_values */
+    background: linear-gradient(to left, #edf7ff, #f9f4ff);
+    padding: $euiSizeXS;
+  }
+
+  &.queryAssist__summary_banner_dark {
+    background: transparent;
+    padding: $euiSizeXS;
+  }
+
   &.queryAssist__callout {
     margin-top: $euiSizeXS;
   }

src/plugins/query_enhancements/public/query_assist/components/query_assist_button.test.tsx

@@ -3,8 +3,8 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import { fireEvent, render, screen } from '@testing-library/react';
-import React from 'react';
-import { QueryAssistButton, QueryAssistButtonProps } from './query_assist_button';
+import React, { ComponentProps } from 'react';
+import { QueryAssistButton } from './query_assist_button';
 import { useQueryAssist } from '../hooks';
 
 jest.mock('../hooks', () => ({
@@ -15,7 +15,7 @@ describe('query assist button', () => {
   const setIsCollapsed = jest.fn();
   const updateIsQueryAssistCollapsed = jest.fn();
 
-  const props: QueryAssistButtonProps = {
+  const props: ComponentProps<typeof QueryAssistButton> = {
     dependencies: {
       isCollapsed: false,
       setIsCollapsed,

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

@@ -3,36 +3,36 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import { EuiButtonIcon, EuiFlexItem } from '@elastic/eui';
-import React from 'react';
+import React, { useCallback } from 'react';
 import { i18n } from '@osd/i18n';
 import { useQueryAssist } from '../hooks';
 import collapsedIcon from '../../assets/sparkle_solid.svg';
-import expandIcon from '../../assets/sparkle_color.svg';
+import expandIcon from '../../assets/sparkle_mark.svg';
 import { QueryEditorExtensionDependencies } from '../../../../data/public';
 
-export interface QueryAssistButtonProps {
+interface QueryAssistButtonProps {
   dependencies: QueryEditorExtensionDependencies;
 }
 
 export const QueryAssistButton: React.FC<QueryAssistButtonProps> = (props) => {
   const { isQueryAssistCollapsed, updateIsQueryAssistCollapsed } = useQueryAssist();
 
-  const onClick = () => {
+  const onClick = useCallback(() => {
     if (props.dependencies.isCollapsed) {
       props.dependencies.setIsCollapsed(false);
       updateIsQueryAssistCollapsed(false);
     } else {
       updateIsQueryAssistCollapsed(!isQueryAssistCollapsed);
     }
-  };
+  }, [props.dependencies, isQueryAssistCollapsed, updateIsQueryAssistCollapsed]);
 
   return (
     <EuiFlexItem grow={false}>
       <EuiButtonIcon
         iconType={
           !props.dependencies.isCollapsed && !isQueryAssistCollapsed ? expandIcon : collapsedIcon
         }
-        aria-label={i18n.translate('discover.queryControls.queryAssistToggle', {
+        aria-label={i18n.translate('queryEnhancements.queryAssist.button.ariaLable', {
           defaultMessage: `Query Assist Toggle`,
         })}
         onClick={onClick}

src/plugins/query_enhancements/public/query_assist/components/query_assist_summary.test.tsx

@@ -3,10 +3,10 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import { fireEvent, render, screen } from '@testing-library/react';
-import React from 'react';
+import React, { ComponentProps } from 'react';
 import { coreMock } from '../../../../../core/public/mocks';
 import { BehaviorSubject } from 'rxjs';
-import { QueryAssistSummary, QueryAssistSummaryProps, QueryContext } from './query_assist_summary';
+import { QueryAssistSummary } from './query_assist_summary';
 import { useQueryAssist } from '../hooks';
 
 jest.mock('react', () => ({
@@ -21,36 +21,35 @@ jest.mock('../hooks', () => ({
 describe('query assist summary', () => {
   const PPL = 'ppl';
   const question = 'Are there any errors in my logs?';
-  const queryContext: QueryContext = {
+  const queryContext = {
     question,
     query: PPL,
     queryResults: [{ size: 1 }],
   };
 
-  const emptyResultQueryContext: QueryContext = {
+  const emptyResultQueryContext = {
     question,
     query: PPL,
     queryResults: [],
   };
 
+  const coreSetupMock = coreMock.createSetup({});
+  const httpMock = coreSetupMock.http;
+  const data = new BehaviorSubject<any[]>([]);
+  const question$ = new BehaviorSubject<string>('');
+  const query$ = new BehaviorSubject<string>('');
   const reportUiStatsMock = jest.fn();
-  const getQuery = jest.fn();
   const setSummary = jest.fn();
   const setLoading = jest.fn();
   const setQueryContext = jest.fn();
   const setFeedback = jest.fn();
   const setIsAssistantEnabledByCapability = jest.fn();
-
-  const coreSetupMock = coreMock.createSetup({});
-  const httpMock = coreSetupMock.http;
-  const data = new BehaviorSubject<any[]>([]);
-  const question$ = new BehaviorSubject<string>('');
-  const query$ = new BehaviorSubject<string>('');
+  const getQuery = jest.fn();
   const dataMock = {
     query: {
       queryString: {
+        getUpdates$: () => query$,
         getQuery,
-        query$,
       },
     },
     search: {
@@ -73,7 +72,7 @@ describe('query assist summary', () => {
       CLICK: 'click',
     },
   };
-  const props: QueryAssistSummaryProps = {
+  const props: ComponentProps<typeof QueryAssistSummary> = {
     data: dataMock,
     http: httpMock,
     usageCollection: usageCollectionMock,
@@ -132,6 +131,7 @@ describe('query assist summary', () => {
       isAssistantEnabledByCapability,
       setIsAssistantEnabledByCapability,
     ]);
+    React.useState.mockImplementationOnce(() => [undefined, jest.fn()]);
     useQueryAssist.mockImplementationOnce(() => ({
       question: 'question',
       question$,
@@ -212,7 +212,7 @@ describe('query assist summary', () => {
     fireEvent.click(screen.getByTestId('queryAssist_summary_buttons_thumbup'));
     expect(setFeedback).toHaveBeenCalledWith(true);
     expect(reportUiStatsMock).toHaveBeenCalledWith(
-      'query-assistant',
+      'query-assist',
       'click',
       expect.stringMatching(/^thumbup/)
     );
@@ -226,7 +226,7 @@ describe('query assist summary', () => {
     fireEvent.click(screen.getByTestId('queryAssist_summary_buttons_thumbdown'));
     expect(setFeedback).toHaveBeenCalledWith(true);
     expect(reportUiStatsMock).toHaveBeenCalledWith(
-      'query-assistant',
+      'query-assist',
       'click',
       expect.stringMatching(/^thumbdown/)
     );
@@ -251,6 +251,9 @@ describe('query assist summary', () => {
         question,
         ppl: PPL,
       }),
+      query: {
+        dataSourceId: undefined,
+      },
     });
     await sleep(2000);
     expect(setSummary).toHaveBeenNthCalledWith(1, null);
@@ -318,4 +321,13 @@ describe('query assist summary', () => {
     await sleep(2000);
     expect(setQueryContext).toHaveBeenCalledTimes(1);
   });
+
+  it('should reset feedback state if re-fetch summary', async () => {
+    mockUseState('summary', LOADING.NO, queryContext, FEEDBACK.YES);
+    const RESPONSE_TEXT = 'response';
+    httpMock.post.mockResolvedValue(RESPONSE_TEXT);
+    renderQueryAssistSummary(COLLAPSED.NO);
+    await sleep(2000);
+    expect(setFeedback).toHaveBeenCalledWith(FEEDBACK.NO);
+  });
 });