fix time range when running query

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

public/components/common/search/__tests__/search.test.tsx

@@ -3,15 +3,18 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { applyMiddleware, createStore } from '@reduxjs/toolkit';
+import { render } from '@testing-library/react';
 import { configure, mount } from 'enzyme';
 import Adapter from 'enzyme-adapter-react-16';
 import React from 'react';
-import { applyMiddleware, createStore } from '@reduxjs/toolkit';
-import { rootReducer } from '../../../../framework/redux/reducers';
 import { Provider } from 'react-redux';
-import { Search } from '../search';
 import thunk from 'redux-thunk';
+import { coreRefs } from '../../../../framework/core_refs';
+import { rootReducer } from '../../../../framework/redux/reducers';
 import { initialTabId } from '../../../../framework/redux/store/shared_state';
+import * as hookExports from '../../../event_analytics/explorer/query_assist/hooks';
+import { Search } from '../search';
 
 describe('Explorer Search component', () => {
   configure({ adapter: new Adapter() });
@@ -27,3 +30,63 @@ describe('Explorer Search component', () => {
     expect(wrapper).toMatchSnapshot();
   });
 });
+
+describe('Search state', () => {
+  const catIndicesSpy = jest.spyOn(hookExports, 'useCatIndices');
+  const getIndexPatternsSpy = jest.spyOn(hookExports, 'useGetIndexPatterns');
+  const store = createStore(rootReducer, applyMiddleware(thunk));
+
+  beforeEach(() => {
+    coreRefs.queryAssistEnabled = true;
+  });
+
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it('selects logs sample data over others', async () => {
+    catIndicesSpy.mockReturnValue({
+      data: [{ label: 'opensearch_dashboards_sample_data_flights' }],
+      loading: false,
+      refresh: jest.fn(),
+    });
+    getIndexPatternsSpy.mockReturnValue({
+      data: [{ label: 'test-index' }, { label: 'opensearch_dashboards_sample_data_logs' }],
+      loading: false,
+      refresh: jest.fn(),
+    });
+    const component = render(
+      <Provider store={store}>
+        <Search
+          tabId={initialTabId}
+          handleQueryChange={jest.fn()}
+          pplService={{ fetch: () => Promise.resolve() }}
+        />
+      </Provider>
+    );
+    expect(component.getByText('opensearch_dashboards_sample_data_logs')).toBeInTheDocument();
+  });
+
+  it('selects other sample data', async () => {
+    catIndicesSpy.mockReturnValue({
+      data: [{ label: 'test-index' }, { label: 'opensearch_dashboards_sample_data_flights' }],
+      loading: false,
+      refresh: jest.fn(),
+    });
+    getIndexPatternsSpy.mockReturnValue({
+      data: [],
+      loading: false,
+      refresh: jest.fn(),
+    });
+    const component = render(
+      <Provider store={store}>
+        <Search
+          tabId={initialTabId}
+          handleQueryChange={jest.fn()}
+          pplService={{ fetch: () => Promise.resolve() }}
+        />
+      </Provider>
+    );
+    expect(component.getByText('opensearch_dashboards_sample_data_flights')).toBeInTheDocument();
+  });
+});

public/components/common/search/search.tsx

@@ -30,7 +30,11 @@ import {
   OLLY_QUERY_ASSISTANT,
   RAW_QUERY,
 } from '../../../../common/constants/explorer';
-import { PPL_SPAN_REGEX } from '../../../../common/constants/shared';
+import {
+  PPL_SPAN_REGEX,
+  QUERY_ASSIST_END_TIME,
+  QUERY_ASSIST_START_TIME,
+} from '../../../../common/constants/shared';
 import { uiSettingsService } from '../../../../common/utils';
 import { useFetchEvents } from '../../../components/event_analytics/hooks';
 import { usePolling } from '../../../components/hooks/use_polling';
@@ -277,7 +281,11 @@ export const Search = (props: any) => {
       dispatch(changeQuery({ tabId, query: { [RAW_QUERY]: tempQuery } }));
     });
     onQuerySearch(queryLang);
-    handleTimePickerChange([startTime, endTime]);
+    if (coreRefs.queryAssistEnabled) {
+      handleTimePickerChange([QUERY_ASSIST_START_TIME, QUERY_ASSIST_END_TIME]);
+    } else {
+      handleTimePickerChange([startTime, endTime]);
+    }
     setNeedsUpdate(false);
   };
 
@@ -299,7 +307,10 @@ export const Search = (props: any) => {
     const sampleLogOption = indicesAndIndexPatterns.find(
       (option) => option.label === 'opensearch_dashboards_sample_data_logs'
     );
-    if (sampleLogOption) setSelectedIndex([sampleLogOption]);
+    if (sampleLogOption) {
+      setSelectedIndex([sampleLogOption]);
+      return;
+    }
     const sampleDataOption = indicesAndIndexPatterns.find((option) =>
       option.label.startsWith('opensearch_dashboards_sample_data_')
     );

public/components/event_analytics/explorer/query_assist/input.tsx

@@ -22,7 +22,10 @@ import React, { useEffect, useRef, useState } from 'react';
 import { useDispatch, useSelector } from 'react-redux';
 import { RAW_QUERY } from '../../../../../common/constants/explorer';
 import { ERROR_DETAILS, QUERY_ASSIST_API } from '../../../../../common/constants/query_assist';
-import { QUERY_ASSIST_START_TIME } from '../../../../../common/constants/shared';
+import {
+  QUERY_ASSIST_END_TIME,
+  QUERY_ASSIST_START_TIME,
+} from '../../../../../common/constants/shared';
 import { getOSDHttp } from '../../../../../common/utils';
 import { coreRefs } from '../../../../framework/core_refs';
 import chatLogo from '../../../datasources/icons/query-assistant-logo.svg';
@@ -35,7 +38,12 @@ import {
 } from '../../redux/slices/query_assistant_summarization_slice';
 import { reset, selectQueryResult } from '../../redux/slices/query_result_slice';
 import { changeQuery, selectQueries } from '../../redux/slices/query_slice';
-import { EmptyIndexCallOut, EmptyQueryCallOut, PPLGeneratedCallOut, ProhibitedQueryCallOut } from './callouts';
+import {
+  EmptyIndexCallOut,
+  EmptyQueryCallOut,
+  PPLGeneratedCallOut,
+  ProhibitedQueryCallOut,
+} from './callouts';
 
 class ProhibitedQueryError extends Error {
   constructor(message?: string) {
@@ -289,7 +297,7 @@ export const QueryAssistInput: React.FC<React.PropsWithChildren<Props>> = (props
       dispatch(setLoading({ tabId: props.tabId, loading: true }));
       dismissCallOut();
       await request();
-      await props.handleTimePickerChange([QUERY_ASSIST_START_TIME, 'now']);
+      await props.handleTimePickerChange([QUERY_ASSIST_START_TIME, QUERY_ASSIST_END_TIME]);
       await props.handleTimeRangePickerRefresh(undefined, true);
     } catch (err) {
       const error = formatError(err);