add more tests for query assist

Signed-off-by: Joshua Li <joshuali925@gmail.com>

.eslintrc.js

@@ -1,7 +1,27 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+const LICENSE_HEADER = `/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */`;
+
 module.exports = {
   root: true,
   extends: ['@elastic/eslint-config-kibana', 'plugin:@elastic/eui/recommended'],
-  rules: {
-    '@osd/eslint/require-license-header': 'off',
-  },
+  overrides: [
+    {
+      files: ['**/*.{js,ts,tsx}'],
+      rules: {
+        '@osd/eslint/require-license-header': [
+          'error',
+          {
+            licenses: [LICENSE_HEADER],
+          },
+        ],
+      },
+    },
+  ],
 };

public/query_assist/components/index_selector.tsx

@@ -31,7 +31,7 @@ export const IndexSelector: React.FC<IndexSelectorProps> = (props) => {
 
   return (
     <EuiComboBox
-      style={{ width: 300 }}
+      style={{ width: 500 }}
       placeholder="Select an index"
       isClearable={false}
       prepend={<EuiText>Index</EuiText>}

public/query_assist/components/query_assist_banner.test.tsx

@@ -0,0 +1,50 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { I18nProvider } from '@osd/i18n/react';
+import { fireEvent, render } from '@testing-library/react';
+import React, { ComponentProps } from 'react';
+import { QueryAssistBanner } from './query_assist_banner';
+
+jest.mock('../../services', () => ({
+  getStorage: () => ({
+    get: jest.fn(),
+    set: jest.fn(),
+  }),
+}));
+
+type QueryAssistBannerProps = ComponentProps<typeof QueryAssistBanner>;
+
+const renderQueryAssistBanner = (overrideProps: Partial<QueryAssistBannerProps> = {}) => {
+  const props: QueryAssistBannerProps = Object.assign<
+    QueryAssistBannerProps,
+    Partial<QueryAssistBannerProps>
+  >(
+    {
+      languages: ['test-lang1', 'test-lang2'],
+    },
+    overrideProps
+  );
+  const component = render(
+    <I18nProvider>
+      <QueryAssistBanner {...props} />
+    </I18nProvider>
+  );
+  return { component, props: props as jest.MockedObjectDeep<QueryAssistBannerProps> };
+};
+
+describe('<QueryAssistBanner /> spec', () => {
+  it('should dismiss callout', async () => {
+    const { component } = renderQueryAssistBanner();
+    expect(
+      component.getByText('Natural Language Query Generation for test-lang1, test-lang2')
+    ).toBeInTheDocument();
+
+    fireEvent.click(component.getByTestId('closeCallOutButton'));
+    expect(
+      component.queryByText('Natural Language Query Generation for test-lang1, test-lang2')
+    ).toBeNull();
+  });
+});

public/query_assist/components/query_assist_input.test.tsx

@@ -0,0 +1,76 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { I18nProvider } from '@osd/i18n/react';
+import { fireEvent, render } from '@testing-library/react';
+import React, { ComponentProps } from 'react';
+import { SuggestionsComponentProps } from '../../../../../src/plugins/data/public/ui/typeahead/suggestions_component';
+import { QueryAssistInput } from './query_assist_input';
+
+jest.mock('../../services', () => ({
+  getData: () => ({
+    ui: {
+      SuggestionsComponent: ({ show, suggestions, onClick }: SuggestionsComponentProps) => (
+        <div data-test-subj="suggestions-component">
+          {show &&
+            suggestions.map((s, i) => (
+              <button key={i} onClick={() => onClick(s)}>
+                {s.text}
+              </button>
+            ))}
+        </div>
+      ),
+    },
+  }),
+}));
+
+const mockPersistedLog = {
+  get: () => ['mock suggestion 1', 'mock suggestion 2'],
+} as any;
+
+type QueryAssistInputProps = ComponentProps<typeof QueryAssistInput>;
+
+const renderQueryAssistInput = (overrideProps: Partial<QueryAssistInputProps> = {}) => {
+  const props: QueryAssistInputProps = Object.assign<
+    QueryAssistInputProps,
+    Partial<QueryAssistInputProps>
+  >(
+    { inputRef: { current: null }, persistedLog: mockPersistedLog, isDisabled: false },
+    overrideProps
+  );
+  const component = render(
+    <I18nProvider>
+      <QueryAssistInput {...props} />
+    </I18nProvider>
+  );
+  return { component, props: props as jest.MockedObjectDeep<QueryAssistInputProps> };
+};
+
+describe('<QueryAssistInput /> spec', () => {
+  it('should display input', () => {
+    const { component } = renderQueryAssistInput();
+    const inputElement = component.getByTestId('query-assist-input-field-text') as HTMLInputElement;
+    expect(inputElement).toBeInTheDocument();
+    fireEvent.change(inputElement, { target: { value: 'new value' } });
+    expect(inputElement.value).toBe('new value');
+  });
+
+  it('should display suggestions on input click', () => {
+    const { component } = renderQueryAssistInput();
+    const inputElement = component.getByTestId('query-assist-input-field-text') as HTMLInputElement;
+    fireEvent.click(inputElement);
+    const suggestionsComponent = component.getByTestId('suggestions-component');
+    expect(suggestionsComponent).toBeInTheDocument();
+  });
+
+  it('should update input value on suggestion click', () => {
+    const { component } = renderQueryAssistInput();
+    const inputElement = component.getByTestId('query-assist-input-field-text') as HTMLInputElement;
+    fireEvent.click(inputElement);
+    const suggestionButton = component.getByText('mock suggestion 1');
+    fireEvent.click(suggestionButton);
+    expect(inputElement.value).toBe('mock suggestion 1');
+  });
+});

public/query_assist/components/query_assist_input.tsx

@@ -74,6 +74,7 @@ export const QueryAssistInput: React.FC<QueryAssistInputProps> = (props) => {
     <EuiOutsideClickDetector onOutsideClick={() => setIsSuggestionsVisible(false)}>
       <div>
         <EuiFieldText
+          data-test-subj="query-assist-input-field-text"
           inputRef={props.inputRef}
           value={value}
           disabled={props.isDisabled}

public/query_assist/hooks/use_generate.test.ts

@@ -0,0 +1,115 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { act, renderHook } from '@testing-library/react-hooks/dom';
+import { coreMock } from '../../../../../src/core/public/mocks';
+import { useOpenSearchDashboards } from '../../../../../src/plugins/opensearch_dashboards_react/public';
+import { useGenerateQuery } from './use_generate';
+
+const coreSetup = coreMock.createSetup();
+const mockHttp = coreSetup.http;
+
+jest.mock('../../../../../src/plugins/opensearch_dashboards_react/public', () => ({
+  useOpenSearchDashboards: jest.fn(),
+  withOpenSearchDashboards: jest.fn((component: React.Component) => component),
+}));
+
+describe('useGenerateQuery', () => {
+  beforeEach(() => {
+    (useOpenSearchDashboards as jest.MockedFunction<typeof useOpenSearchDashboards>)
+      // @ts-ignore for this test we only need http implemented
+      .mockImplementation(() => ({
+        services: {
+          http: mockHttp,
+        },
+      }));
+  });
+
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it.skip('should show loading', async () => {
+    mockHttp.post.mockResolvedValueOnce({ query: 'test query' });
+    const { result, waitForNextUpdate } = renderHook(() => useGenerateQuery());
+    const { generateQuery } = result.current;
+
+    await act(async () => {
+      generateQuery({
+        index: 'test',
+        language: 'test-lang',
+        question: 'test question',
+      });
+
+      // Assert that loading is true immediately after generating the query
+      expect(result.current.loading).toBe(true);
+
+      // Wait for the next state update
+      await waitForNextUpdate();
+    });
+
+    // Assert that loading is false after the state update
+    expect(result.current.loading).toBe(false);
+  });
+
+  it('should generate results', async () => {
+    mockHttp.post.mockResolvedValueOnce({ query: 'test query' });
+    const { result } = renderHook(() => useGenerateQuery());
+    const { generateQuery } = result.current;
+
+    await act(async () => {
+      const response = await generateQuery({
+        index: 'test',
+        language: 'test-lang',
+        question: 'test question',
+      });
+
+      expect(response).toEqual({ response: { query: 'test query' } });
+    });
+  });
+
+  it('should handle errors', async () => {
+    const { result } = renderHook(() => useGenerateQuery());
+    const { generateQuery } = result.current;
+    const mockError = new Error('mockError');
+    mockHttp.post.mockRejectedValueOnce(mockError);
+
+    await act(async () => {
+      const response = await generateQuery({
+        index: 'test',
+        language: 'test-lang',
+        question: 'test question',
+      });
+
+      expect(response).toEqual({ error: mockError });
+      expect(result.current.loading).toBe(false);
+    });
+  });
+
+  it('should abort previous call', async () => {
+    const { result } = renderHook(() => useGenerateQuery());
+    const { generateQuery, abortControllerRef } = result.current;
+
+    await act(async () => {
+      await generateQuery({ index: 'test', language: 'test-lang', question: 'test question' });
+      const controller = abortControllerRef.current;
+      await generateQuery({ index: 'test', language: 'test-lang', question: 'test question' });
+
+      expect(controller?.signal.aborted).toBe(true);
+    });
+  });
+
+  it('should abort call with controller', async () => {
+    const { result } = renderHook(() => useGenerateQuery());
+    const { generateQuery, abortControllerRef } = result.current;
+
+    await act(async () => {
+      await generateQuery({ index: 'test', language: 'test-lang', question: 'test question' });
+      abortControllerRef.current?.abort();
+
+      expect(abortControllerRef.current?.signal.aborted).toBe(true);
+    });
+  });
+});

public/query_assist/utils/create_extension.test.tsx

@@ -0,0 +1,100 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { act, render, screen } from '@testing-library/react';
+import React from 'react';
+import { coreMock } from '../../../../../src/core/public/mocks';
+import { IIndexPattern } from '../../../../../src/plugins/data/public';
+import { PublicConfig } from '../../plugin';
+import { createQueryAssistExtension } from './create_extension';
+
+const coreSetupMock = coreMock.createSetup();
+const httpMock = coreSetupMock.http;
+
+jest.mock('../../services', () => ({
+  getData: jest.fn().mockReturnValue({
+    indexPatterns: {
+      get: jest.fn().mockResolvedValue({ id: 'test-pattern' }),
+    },
+  }),
+}));
+
+jest.mock('.', () => ({
+  getMdsDataSourceId: jest.fn().mockResolvedValue('mock-data-source-id'),
+}));
+
+jest.mock('../components', () => ({
+  QueryAssistBar: jest.fn(() => <div>QueryAssistBar</div>),
+}));
+
+jest.mock('../components/query_assist_banner', () => ({
+  QueryAssistBanner: jest.fn(() => <div>QueryAssistBanner</div>),
+}));
+
+describe('CreateExtension', () => {
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  const config: PublicConfig = {
+    queryAssist: {
+      supportedLanguages: [{ language: 'PPL', agentConfig: 'os_query_assist_ppl' }],
+    },
+  };
+
+  it('should be enabled if at least one language is configured', async () => {
+    httpMock.get.mockResolvedValueOnce({ configuredLanguages: ['PPL'] });
+    const extension = createQueryAssistExtension(httpMock, config);
+    const isEnabled = await extension.isEnabled({ language: 'PPL' });
+    expect(isEnabled).toBeTruthy();
+    expect(httpMock.get).toBeCalledWith('/api/enhancements/assist/languages', {
+      query: { dataSourceId: 'mock-data-source-id' },
+    });
+  });
+
+  it('should be disabled for unsupported language', async () => {
+    httpMock.get.mockRejectedValueOnce(new Error('network failure'));
+    const extension = createQueryAssistExtension(httpMock, config);
+    const isEnabled = await extension.isEnabled({ language: 'PPL' });
+    expect(isEnabled).toBeFalsy();
+    expect(httpMock.get).toBeCalledWith('/api/enhancements/assist/languages', {
+      query: { dataSourceId: 'mock-data-source-id' },
+    });
+  });
+
+  it('should render the component if language is supported', async () => {
+    httpMock.get.mockResolvedValueOnce({ configuredLanguages: ['PPL'] });
+    const extension = createQueryAssistExtension(httpMock, config);
+    const component = extension.getComponent?.({
+      language: 'PPL',
+      indexPatterns: [{ id: 'test-pattern' }] as IIndexPattern[],
+    });
+
+    if (!component) throw new Error('QueryEditorExtensions Component is undefined');
+
+    await act(async () => {
+      render(component);
+    });
+
+    expect(screen.getByText('QueryAssistBar')).toBeInTheDocument();
+  });
+
+  it('should render the banner if language is not supported', async () => {
+    httpMock.get.mockResolvedValueOnce({ configuredLanguages: ['PPL'] });
+    const extension = createQueryAssistExtension(httpMock, config);
+    const banner = extension.getBanner?.({
+      language: 'DQL',
+      indexPatterns: [{ id: 'test-pattern' }] as IIndexPattern[],
+    });
+
+    if (!banner) throw new Error('QueryEditorExtensions Banner is undefined');
+
+    await act(async () => {
+      render(banner);
+    });
+
+    expect(screen.getByText('QueryAssistBanner')).toBeInTheDocument();
+  });
+});