[ES|QL] Adds more tests in unified histogram component (#165568)

## Summary

Closes https://github.com/elastic/kibana/issues/165431

Adds more unit tests in unified_histogram plugin, as a follow up from
ES|QL

### Checklist

- [x] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios

---------

Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>

src/plugins/discover/public/application/main/hooks/use_test_based_query_language.test.tsx

@@ -142,6 +142,7 @@ describe('useTextBasedQueryLanguage', () => {
           flattened: { field1: 1 },
         } as unknown as DataTableRecord,
       ],
+      // transformational command
       query: { esql: 'from the-data-view-title | keep field1' },
     });
     await waitFor(() => expect(replaceUrlState).toHaveBeenCalledTimes(1));
@@ -153,6 +154,35 @@ describe('useTextBasedQueryLanguage', () => {
       });
     });
   });
+
+  test('changing a text based query with no transformational commands should only change dataview state when loading and finished', async () => {
+    const { replaceUrlState, stateContainer } = renderHookWithContext(false);
+    const documents$ = stateContainer.dataState.data$.documents$;
+    stateContainer.dataState.data$.documents$.next(msgComplete);
+    await waitFor(() => expect(replaceUrlState).toHaveBeenCalledTimes(1));
+    replaceUrlState.mockReset();
+
+    documents$.next({
+      recordRawType: RecordRawType.PLAIN,
+      fetchStatus: FetchStatus.PARTIAL,
+      result: [
+        {
+          id: '1',
+          raw: { field1: 1 },
+          flattened: { field1: 1 },
+        } as unknown as DataTableRecord,
+      ],
+      // non transformational command
+      query: { esql: 'from the-data-view-title | where field1 > 0' },
+    });
+    await waitFor(() => expect(replaceUrlState).toHaveBeenCalledTimes(1));
+
+    await waitFor(() => {
+      expect(replaceUrlState).toHaveBeenCalledWith({
+        index: 'the-data-view-id',
+      });
+    });
+  });
   test('only changing a text based query with same result columns should not change columns', async () => {
     const { replaceUrlState, stateContainer } = renderHookWithContext(false);
 
@@ -271,6 +301,65 @@ describe('useTextBasedQueryLanguage', () => {
     });
   });
 
+  test('it should not overwrite existing state columns on initial fetch and non transformational commands', async () => {
+    const { replaceUrlState, stateContainer } = renderHookWithContext(false, {
+      columns: ['field1'],
+      index: 'the-data-view-id',
+    });
+    const documents$ = stateContainer.dataState.data$.documents$;
+
+    documents$.next({
+      recordRawType: RecordRawType.PLAIN,
+      fetchStatus: FetchStatus.PARTIAL,
+      result: [
+        {
+          id: '1',
+          raw: { field1: 1, field2: 2 },
+          flattened: { field1: 1 },
+        } as unknown as DataTableRecord,
+      ],
+      query: { esql: 'from the-data-view-title | WHERE field2=1' },
+    });
+    expect(replaceUrlState).toHaveBeenCalledTimes(0);
+  });
+
+  test('it should overwrite existing state columns on transitioning from a query with non transformational commands to a query with transformational', async () => {
+    const { replaceUrlState, stateContainer } = renderHookWithContext(false, {
+      index: 'the-data-view-id',
+    });
+    const documents$ = stateContainer.dataState.data$.documents$;
+
+    documents$.next({
+      recordRawType: RecordRawType.PLAIN,
+      fetchStatus: FetchStatus.PARTIAL,
+      result: [
+        {
+          id: '1',
+          raw: { field1: 1, field2: 2 },
+          flattened: { field1: 1 },
+        } as unknown as DataTableRecord,
+      ],
+      query: { esql: 'from the-data-view-title | WHERE field2=1' },
+    });
+    expect(replaceUrlState).toHaveBeenCalledTimes(0);
+    documents$.next({
+      recordRawType: RecordRawType.PLAIN,
+      fetchStatus: FetchStatus.PARTIAL,
+      result: [
+        {
+          id: '1',
+          raw: { field1: 1 },
+          flattened: { field1: 1 },
+        } as unknown as DataTableRecord,
+      ],
+      query: { esql: 'from the-data-view-title | keep field1' },
+    });
+    await waitFor(() => expect(replaceUrlState).toHaveBeenCalledTimes(1));
+    expect(replaceUrlState).toHaveBeenCalledWith({
+      columns: ['field1'],
+    });
+  });
+
   test('it should not overwrite state column when successfully fetching after an error fetch', async () => {
     const { replaceUrlState, stateContainer } = renderHookWithContext(false, {
       columns: [],

src/plugins/unified_histogram/public/layout/hooks/use_lens_suggestions.test.ts

@@ -0,0 +1,178 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+import { renderHook } from '@testing-library/react-hooks';
+import { dataPluginMock } from '@kbn/data-plugin/public/mocks';
+import { calculateBounds } from '@kbn/data-plugin/public';
+import { deepMockedFields, buildDataViewMock } from '@kbn/discover-utils/src/__mocks__';
+import { allSuggestionsMock } from '../../__mocks__/suggestions';
+import { useLensSuggestions } from './use_lens_suggestions';
+
+describe('useLensSuggestions', () => {
+  const dataMock = dataPluginMock.createStartContract();
+  dataMock.query.timefilter.timefilter.calculateBounds = (timeRange) => {
+    return calculateBounds(timeRange);
+  };
+  const dataViewMock = buildDataViewMock({
+    name: 'the-data-view',
+    fields: deepMockedFields,
+    timeFieldName: '@timestamp',
+  });
+
+  test('should return empty suggestions for non aggregate query', async () => {
+    const { result } = renderHook(() => {
+      return useLensSuggestions({
+        dataView: dataViewMock,
+        query: undefined,
+        isPlainRecord: false,
+        data: dataMock,
+        lensSuggestionsApi: jest.fn(),
+      });
+    });
+    const current = result.current;
+    expect(current).toStrictEqual({
+      allSuggestions: [],
+      currentSuggestion: undefined,
+      isOnHistogramMode: false,
+      suggestionUnsupported: false,
+    });
+  });
+
+  test('should return suggestions for aggregate query', async () => {
+    const { result } = renderHook(() => {
+      return useLensSuggestions({
+        dataView: dataViewMock,
+        query: { esql: 'from the-data-view | stats maxB = max(bytes)' },
+        isPlainRecord: true,
+        columns: [
+          {
+            id: 'var0',
+            name: 'var0',
+            meta: {
+              type: 'number',
+            },
+          },
+        ],
+        data: dataMock,
+        lensSuggestionsApi: jest.fn(() => allSuggestionsMock),
+      });
+    });
+    const current = result.current;
+    expect(current).toStrictEqual({
+      allSuggestions: allSuggestionsMock,
+      currentSuggestion: allSuggestionsMock[0],
+      isOnHistogramMode: false,
+      suggestionUnsupported: false,
+    });
+  });
+
+  test('should return suggestionUnsupported if no timerange is provided and no suggestions returned by the api', async () => {
+    const { result } = renderHook(() => {
+      return useLensSuggestions({
+        dataView: dataViewMock,
+        query: { esql: 'from the-data-view | stats maxB = max(bytes)' },
+        isPlainRecord: true,
+        columns: [
+          {
+            id: 'var0',
+            name: 'var0',
+            meta: {
+              type: 'number',
+            },
+          },
+        ],
+        data: dataMock,
+        lensSuggestionsApi: jest.fn(),
+      });
+    });
+    const current = result.current;
+    expect(current).toStrictEqual({
+      allSuggestions: [],
+      currentSuggestion: undefined,
+      isOnHistogramMode: false,
+      suggestionUnsupported: true,
+    });
+  });
+
+  test('should return histogramSuggestion if no suggestions returned by the api', async () => {
+    const firstMockReturn = undefined;
+    const secondMockReturn = allSuggestionsMock;
+    const lensSuggestionsApi = jest
+      .fn()
+      .mockReturnValueOnce(firstMockReturn) // will return to firstMockReturn object firstly
+      .mockReturnValueOnce(secondMockReturn); // will return to secondMockReturn object secondly
+
+    const { result } = renderHook(() => {
+      return useLensSuggestions({
+        dataView: dataViewMock,
+        query: { esql: 'from the-data-view | limit 100' },
+        isPlainRecord: true,
+        columns: [
+          {
+            id: 'var0',
+            name: 'var0',
+            meta: {
+              type: 'number',
+            },
+          },
+        ],
+        data: dataMock,
+        lensSuggestionsApi,
+        timeRange: {
+          from: '2023-09-03T08:00:00.000Z',
+          to: '2023-09-04T08:56:28.274Z',
+        },
+      });
+    });
+    const current = result.current;
+    expect(current).toStrictEqual({
+      allSuggestions: [],
+      currentSuggestion: allSuggestionsMock[0],
+      isOnHistogramMode: true,
+      suggestionUnsupported: false,
+    });
+  });
+
+  test('should not return histogramSuggestion if no suggestions returned by the api and transformational commands', async () => {
+    const firstMockReturn = undefined;
+    const secondMockReturn = allSuggestionsMock;
+    const lensSuggestionsApi = jest
+      .fn()
+      .mockReturnValueOnce(firstMockReturn) // will return to firstMockReturn object firstly
+      .mockReturnValueOnce(secondMockReturn); // will return to secondMockReturn object secondly
+
+    const { result } = renderHook(() => {
+      return useLensSuggestions({
+        dataView: dataViewMock,
+        query: { esql: 'from the-data-view | limit 100 | keep @timestamp' },
+        isPlainRecord: true,
+        columns: [
+          {
+            id: 'var0',
+            name: 'var0',
+            meta: {
+              type: 'number',
+            },
+          },
+        ],
+        data: dataMock,
+        lensSuggestionsApi,
+        timeRange: {
+          from: '2023-09-03T08:00:00.000Z',
+          to: '2023-09-04T08:56:28.274Z',
+        },
+      });
+    });
+    const current = result.current;
+    expect(current).toStrictEqual({
+      allSuggestions: [],
+      currentSuggestion: undefined,
+      isOnHistogramMode: false,
+      suggestionUnsupported: true,
+    });
+  });
+});

src/plugins/unified_histogram/tsconfig.json

@@ -24,6 +24,7 @@
     "@kbn/ui-actions-plugin",
     "@kbn/kibana-utils-plugin",
     "@kbn/visualizations-plugin",
+    "@kbn/discover-utils",
   ],
   "exclude": [
     "target/**/*",