Tempo: Update status operators (#78730)

* Set keyword operators

* Tests for keyword, string, number option types

public/app/plugins/datasource/tempo/SearchTraceQLEditor/SearchField.test.tsx

@@ -1,40 +1,18 @@
-import { render, screen } from '@testing-library/react';
+import { render, screen, waitFor } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import React from 'react';
 import { initTemplateSrv } from 'test/helpers/initTemplateSrv';
 
+import { LanguageProvider } from '@grafana/data';
 import { FetchError, setTemplateSrv } from '@grafana/runtime';
 
 import { TraceqlFilter, TraceqlSearchScope } from '../dataquery.gen';
 import { TempoDatasource } from '../datasource';
+import TempoLanguageProvider from '../language_provider';
+import { keywordOperators, numberOperators, operators, stringOperators } from '../traceql/traceql';
 
 import SearchField from './SearchField';
 
-const getOptionsV2 = jest.fn().mockImplementation(() => {
-  return new Promise((resolve) => {
-    setTimeout(() => {
-      resolve([
-        {
-          value: 'customer',
-          label: 'customer',
-          type: 'string',
-        },
-        {
-          value: 'driver',
-          label: 'driver',
-          type: 'string',
-        },
-      ]);
-    }, 1000);
-  });
-});
-
-jest.mock('../language_provider', () => {
-  return jest.fn().mockImplementation(() => {
-    return { getOptionsV2 };
-  });
-});
-
 describe('SearchField', () => {
   let templateSrv = initTemplateSrv('key', [{ name: 'templateVariable1' }, { name: 'templateVariable2' }]);
   let user: ReturnType<typeof userEvent.setup>;
@@ -51,7 +29,7 @@ describe('SearchField', () => {
     jest.useRealTimers();
   });
 
-  it('should not render tag if hideTag is true', () => {
+  it('should not render tag if hideTag is true', async () => {
     const updateFilter = jest.fn((val) => {
       return val;
     });
@@ -59,10 +37,12 @@ describe('SearchField', () => {
 
     const { container } = renderSearchField(updateFilter, filter, [], true);
 
+    await waitFor(async () => {
       expect(container.querySelector(`input[aria-label="select test1 tag"]`)).not.toBeInTheDocument();
       expect(container.querySelector(`input[aria-label="select test1 operator"]`)).toBeInTheDocument();
       expect(container.querySelector(`input[aria-label="select test1 value"]`)).toBeInTheDocument();
     });
+  });
 
   it('should update operator when new value is selected in operator input', async () => {
     const updateFilter = jest.fn((val) => {
@@ -71,7 +51,7 @@ describe('SearchField', () => {
     const filter: TraceqlFilter = { id: 'test1', operator: '=', valueType: 'string', tag: 'test-tag' };
     const { container } = renderSearchField(updateFilter, filter);
 
-    const select = await container.querySelector(`input[aria-label="select test1 operator"]`);
+    const select = container.querySelector(`input[aria-label="select test1 operator"]`);
     expect(select).not.toBeNull();
     expect(select).toBeInTheDocument();
     if (select) {
@@ -95,7 +75,7 @@ describe('SearchField', () => {
     };
     const { container } = renderSearchField(updateFilter, filter);
 
-    const select = await container.querySelector(`input[aria-label="select test1 value"]`);
+    const select = container.querySelector(`input[aria-label="select test1 value"]`);
     expect(select).not.toBeNull();
     expect(select).toBeInTheDocument();
     if (select) {
@@ -178,14 +158,112 @@ describe('SearchField', () => {
       expect(await screen.findByText('$templateVariable2')).toBeInTheDocument();
     }
   });
+
+  it('should only show keyword operators if options tag type is keyword', async () => {
+    const filter: TraceqlFilter = { id: 'test1', operator: '=', valueType: 'string', tag: 'test-tag' };
+    const lp = {
+      getOptionsV2: jest.fn().mockReturnValue([
+        {
+          value: 'ok',
+          label: 'ok',
+          type: 'keyword',
+        },
+      ]),
+    } as unknown as TempoLanguageProvider;
+
+    const { container } = renderSearchField(jest.fn(), filter, [], false, lp);
+    const select = container.querySelector(`input[aria-label="select test1 operator"]`);
+    if (select) {
+      await user.click(select);
+      await waitFor(async () => {
+        expect(screen.getByText('Equals')).toBeInTheDocument();
+        expect(screen.getByText('Not equals')).toBeInTheDocument();
+        operators
+          .filter((op) => !keywordOperators.includes(op))
+          .forEach((op) => {
+            expect(screen.queryByText(op)).not.toBeInTheDocument();
+          });
+      });
+    }
+  });
+
+  it('should only show string operators if options tag type is string', async () => {
+    const filter: TraceqlFilter = { id: 'test1', operator: '=', valueType: 'string', tag: 'test-tag' };
+    const { container } = renderSearchField(jest.fn(), filter);
+    const select = container.querySelector(`input[aria-label="select test1 operator"]`);
+    if (select) {
+      await user.click(select);
+      await waitFor(async () => {
+        expect(screen.getByText('Equals')).toBeInTheDocument();
+        expect(screen.getByText('Not equals')).toBeInTheDocument();
+        expect(screen.getByText('Matches regex')).toBeInTheDocument();
+        expect(screen.getByText('Does not match regex')).toBeInTheDocument();
+        operators
+          .filter((op) => !stringOperators.includes(op))
+          .forEach((op) => {
+            expect(screen.queryByText(op)).not.toBeInTheDocument();
+          });
+      });
+    }
+  });
+
+  it('should only show number operators if options tag type is number', async () => {
+    const filter: TraceqlFilter = { id: 'test1', operator: '=', valueType: 'string', tag: 'test-tag' };
+    const lp = {
+      getOptionsV2: jest.fn().mockReturnValue([
+        {
+          value: 200,
+          label: 200,
+          type: 'int',
+        },
+      ]),
+    } as unknown as TempoLanguageProvider;
+
+    const { container } = renderSearchField(jest.fn(), filter, [], false, lp);
+    const select = container.querySelector(`input[aria-label="select test1 operator"]`);
+    if (select) {
+      await user.click(select);
+      await waitFor(async () => {
+        expect(screen.getByText('Equals')).toBeInTheDocument();
+        expect(screen.getByText('Not equals')).toBeInTheDocument();
+        expect(screen.getByText('Greater')).toBeInTheDocument();
+        expect(screen.getByText('Less')).toBeInTheDocument();
+        expect(screen.getByText('Greater or Equal')).toBeInTheDocument();
+        expect(screen.getByText('Less or Equal')).toBeInTheDocument();
+        operators
+          .filter((op) => !numberOperators.includes(op))
+          .forEach((op) => {
+            expect(screen.queryByText(op)).not.toBeInTheDocument();
+          });
+      });
+    }
+  });
 });
 
 const renderSearchField = (
   updateFilter: (f: TraceqlFilter) => void,
   filter: TraceqlFilter,
   tags?: string[],
-  hideTag?: boolean
+  hideTag?: boolean,
+  lp?: LanguageProvider
 ) => {
+  const languageProvider =
+    lp ||
+    ({
+      getOptionsV2: jest.fn().mockReturnValue([
+        {
+          value: 'customer',
+          label: 'customer',
+          type: 'string',
+        },
+        {
+          value: 'driver',
+          label: 'driver',
+          type: 'string',
+        },
+      ]),
+    } as unknown as TempoLanguageProvider);
+
   const datasource: TempoDatasource = {
     search: {
       filters: [
@@ -198,7 +276,9 @@ const renderSearchField = (
         { id: 'span-name', type: 'static', tag: 'name', operator: '=', scope: TraceqlSearchScope.Span },
       ],
     },
+    languageProvider,
   } as TempoDatasource;
+
   return render(
     <SearchField
       datasource={datasource}

public/app/plugins/datasource/tempo/SearchTraceQLEditor/SearchField.tsx

@@ -13,8 +13,7 @@ import { notifyApp } from '../../../../core/reducers/appNotification';
 import { dispatch } from '../../../../store/store';
 import { TraceqlFilter, TraceqlSearchScope } from '../dataquery.gen';
 import { TempoDatasource } from '../datasource';
-import TempoLanguageProvider from '../language_provider';
+import { operators as allOperators, stringOperators, numberOperators, keywordOperators } from '../traceql/traceql';
-import { operators as allOperators, stringOperators, numberOperators } from '../traceql/traceql';
 
 import { filterScopedTag, operatorSelectableValue } from './utils';
 
@@ -53,7 +52,6 @@ const SearchField = ({
   query,
 }: Props) => {
   const styles = useStyles2(getStyles);
-  const languageProvider = useMemo(() => new TempoLanguageProvider(datasource), [datasource]);
   const scopedTag = useMemo(() => filterScopedTag(filter), [filter]);
   // We automatically change the operator to the regex op when users select 2 or more values
   // However, they expect this to be automatically rolled back to the previous operator once
@@ -63,7 +61,7 @@ const SearchField = ({
 
   const updateOptions = async () => {
     try {
-      return filter.tag ? await languageProvider.getOptionsV2(scopedTag, query) : [];
+      return filter.tag ? await datasource.languageProvider.getOptionsV2(scopedTag, query) : [];
     } catch (error) {
       // Display message if Tempo is connected but search 404's
       if (isFetchError(error) && error?.status === 404) {
@@ -77,7 +75,7 @@ const SearchField = ({
 
   const { loading: isLoadingValues, value: options } = useAsync(updateOptions, [
     scopedTag,
-    languageProvider,
+    datasource.languageProvider,
     setError,
     query,
   ]);
@@ -115,6 +113,9 @@ const SearchField = ({
   const uniqueOptionType = options?.length === optionsOfFirstType?.length ? options?.[0]?.type : undefined;
   let operatorList = allOperators;
   switch (uniqueOptionType) {
+    case 'keyword':
+      operatorList = keywordOperators;
+      break;
     case 'string':
       operatorList = stringOperators;
       break;

public/app/plugins/datasource/tempo/traceql/traceql.ts

@@ -24,6 +24,7 @@ export const languageConfiguration: languages.LanguageConfiguration = {
 };
 
 export const operators = ['=', '!=', '>', '<', '>=', '<=', '=~', '!~'];
+export const keywordOperators = ['=', '!='];
 export const stringOperators = ['=', '!=', '=~', '!~'];
 export const numberOperators = ['=', '!=', '>', '<', '>=', '<='];