Tempo: Tidy up and organise (#90649)

* Move test files to folder

* Update paths

* Tidy up

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

@@ -5,7 +5,7 @@ import { useState } from 'react';
 import { TraceqlSearchScope } from '../dataquery.gen';
 import { TempoDatasource } from '../datasource';
 import TempoLanguageProvider from '../language_provider';
-import { initTemplateSrv } from '../test_utils';
+import { initTemplateSrv } from '../test/test_utils';
 import { TempoQuery } from '../types';
 
 import { GroupByField } from './GroupByField';

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

@@ -6,7 +6,7 @@ import { LanguageProvider } from '@grafana/data';
 import { TraceqlFilter, TraceqlSearchScope } from '../dataquery.gen';
 import { TempoDatasource } from '../datasource';
 import TempoLanguageProvider from '../language_provider';
-import { initTemplateSrv } from '../test_utils';
+import { initTemplateSrv } from '../test/test_utils';
 import { keywordOperators, numberOperators, operators, stringOperators } from '../traceql/traceql';
 
 import SearchField from './SearchField';

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

@@ -4,7 +4,7 @@ import userEvent from '@testing-library/user-event';
 import { TraceqlFilter, TraceqlSearchScope } from '../dataquery.gen';
 import { TempoDatasource } from '../datasource';
 import TempoLanguageProvider from '../language_provider';
-import { initTemplateSrv } from '../test_utils';
+import { initTemplateSrv } from '../test/test_utils';
 import { Scope } from '../types';
 
 import TagsInput from './TagsInput';

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

@@ -7,7 +7,7 @@ import { config } from '@grafana/runtime';
 import { TraceqlSearchScope } from '../dataquery.gen';
 import { TempoDatasource } from '../datasource';
 import TempoLanguageProvider from '../language_provider';
-import { initTemplateSrv } from '../test_utils';
+import { initTemplateSrv } from '../test/test_utils';
 import { TempoQuery } from '../types';
 
 import TraceQLSearch from './TraceQLSearch';

public/app/plugins/datasource/tempo/VariableQueryEditor.test.tsx

@@ -10,7 +10,7 @@ import {
   TempoVariableQueryType,
 } from './VariableQueryEditor';
 import { selectOptionInTest } from './_importedDependencies/test/helpers/selectOptionInTest';
-import { createTempoDatasource } from './mocks';
+import { createTempoDatasource } from './test/mocks';
 
 const refId = 'TempoDatasourceVariableQueryEditor-VariableQuery';
 

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

@@ -41,10 +41,10 @@ import {
   getFieldConfig,
   getEscapedSpanNames,
 } from './datasource';
-import mockJson from './mockJsonResponse.json';
-import mockServiceGraph from './mockServiceGraph.json';
-import { createMetadataRequest, createTempoDatasource } from './mocks';
-import { initTemplateSrv } from './test_utils';
+import mockJson from './test/mockJsonResponse.json';
+import mockServiceGraph from './test/mockServiceGraph.json';
+import { createMetadataRequest, createTempoDatasource } from './test/mocks';
+import { initTemplateSrv } from './test/test_utils';
 import { TempoJsonData, TempoQuery } from './types';
 
 let mockObservable: () => Observable<unknown>;
@@ -68,7 +68,7 @@ describe('Tempo data source', () => {
   describe('runs correctly', () => {
     config.featureToggles.traceQLStreaming = true;
     jest.spyOn(TempoDatasource.prototype, 'isFeatureAvailable').mockImplementation(() => true);
-    const handleStreamingSearch = jest.spyOn(TempoDatasource.prototype, 'handleStreamingSearch');
+    const handleStreamingQuery = jest.spyOn(TempoDatasource.prototype, 'handleStreamingQuery');
     const request = jest.spyOn(TempoDatasource.prototype, '_request');
     const templateSrv: TemplateSrv = { replace: (s: string) => s } as unknown as TemplateSrv;
 
@@ -104,7 +104,7 @@ describe('Tempo data source', () => {
       config.liveEnabled = true;
       const ds = new TempoDatasource(defaultSettings, templateSrv);
       await lastValueFrom(ds.query(traceqlQuery as DataQueryRequest<TempoQuery>));
-      expect(handleStreamingSearch).toHaveBeenCalledTimes(1);
+      expect(handleStreamingQuery).toHaveBeenCalledTimes(1);
       expect(request).toHaveBeenCalledTimes(0);
     });
 
@@ -112,7 +112,7 @@ describe('Tempo data source', () => {
       config.liveEnabled = true;
       const ds = new TempoDatasource(defaultSettings, templateSrv);
       await lastValueFrom(ds.query(traceqlSearchQuery as DataQueryRequest<TempoQuery>));
-      expect(handleStreamingSearch).toHaveBeenCalledTimes(1);
+      expect(handleStreamingQuery).toHaveBeenCalledTimes(1);
       expect(request).toHaveBeenCalledTimes(0);
     });
 
@@ -120,7 +120,7 @@ describe('Tempo data source', () => {
       config.liveEnabled = false;
       const ds = new TempoDatasource(defaultSettings, templateSrv);
       await lastValueFrom(ds.query(traceqlQuery as DataQueryRequest<TempoQuery>));
-      expect(handleStreamingSearch).toHaveBeenCalledTimes(1);
+      expect(handleStreamingQuery).toHaveBeenCalledTimes(1);
       expect(request).toHaveBeenCalledTimes(1);
     });
 
@@ -128,7 +128,7 @@ describe('Tempo data source', () => {
       config.liveEnabled = false;
       const ds = new TempoDatasource(defaultSettings, templateSrv);
       await lastValueFrom(ds.query(traceqlSearchQuery as DataQueryRequest<TempoQuery>));
-      expect(handleStreamingSearch).toHaveBeenCalledTimes(1);
+      expect(handleStreamingQuery).toHaveBeenCalledTimes(1);
       expect(request).toHaveBeenCalledTimes(1);
     });
   });
@@ -365,14 +365,14 @@ describe('Tempo data source', () => {
   describe('test the testDatasource function', () => {
     it('should return a success msg if response.ok is true', async () => {
       mockObservable = () => of({ ok: true });
-      const handleStreamingSearch = jest
-        .spyOn(TempoDatasource.prototype, 'handleStreamingSearch')
+      const handleStreamingQuery = jest
+        .spyOn(TempoDatasource.prototype, 'handleStreamingQuery')
         .mockImplementation(() => of({ data: [] }));
 
       const ds = new TempoDatasource(defaultSettings);
       const response = await ds.testDatasource();
       expect(response.status).toBe('success');
-      expect(handleStreamingSearch).toHaveBeenCalled();
+      expect(handleStreamingQuery).toHaveBeenCalled();
     });
   });
 
@@ -397,7 +397,7 @@ describe('Tempo data source', () => {
       raw: { from: 'now-15m', to: 'now' },
     };
 
-    const request = ds.traceIdQueryRequest(
+    const request = ds.makeTraceIdRequest(
       {
         requestId: 'test',
         interval: '',
@@ -426,7 +426,7 @@ describe('Tempo data source', () => {
       jsonData: { traceQuery: { timeShiftEnabled: false, spanStartTimeShift: '2m', spanEndTimeShift: '4m' } },
     });
 
-    const request = ds.traceIdQueryRequest(
+    const request = ds.makeTraceIdRequest(
       {
         requestId: 'test',
         interval: '',

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

@@ -283,6 +283,19 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
     );
   }
 
+  isTraceQlMetricsQuery(query: string): boolean {
+    // Check whether this is a metrics query by checking if it contains a metrics function
+    const metricsFnRegex =
+      /\|\s*(rate|count_over_time|avg_over_time|max_over_time|min_over_time|quantile_over_time|histogram_over_time|compare)\s*\(/;
+    return !!query.trim().match(metricsFnRegex);
+  }
+
+  isTraceIdQuery(query: string): boolean {
+    const hexOnlyRegex = /^[0-9A-Fa-f]*$/;
+    // Check whether this is a trace ID or traceQL query by checking if it only contains hex characters
+    return !!query.trim().match(hexOnlyRegex);
+  }
+
   query(options: DataQueryRequest<TempoQuery>): Observable<DataQueryResponse> {
     const subQueries: Array<Observable<DataQueryResponse>> = [];
     const filteredTargets = options.targets.filter((target) => !target.hide);
@@ -358,7 +371,7 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
           if (target) {
             const appliedQuery = this.applyVariables(target, options.scopedVars);
             subQueries.push(
-              this.handleMetricsSummary(appliedQuery, generateQueryFromFilters(appliedQuery.filters), options)
+              this.handleMetricsSummaryQuery(appliedQuery, generateQueryFromFilters(appliedQuery.filters), options)
             );
           }
         }
@@ -379,7 +392,7 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
           });
 
           if (this.isStreamingSearchEnabled()) {
-            subQueries.push(this.handleStreamingSearch(options, traceqlSearchTargets, queryValueFromFilters));
+            subQueries.push(this.handleStreamingQuery(options, traceqlSearchTargets, queryValueFromFilters));
           } else {
             subQueries.push(
               this._request('/api/search', {
@@ -459,19 +472,6 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
     return merge(...subQueries);
   }
 
-  isTraceQlMetricsQuery(query: string): boolean {
-    // Check whether this is a metrics query by checking if it contains a metrics function
-    const metricsFnRegex =
-      /\|\s*(rate|count_over_time|avg_over_time|max_over_time|min_over_time|quantile_over_time|histogram_over_time|compare)\s*\(/;
-    return !!query.trim().match(metricsFnRegex);
-  }
-
-  isTraceIdQuery(query: string): boolean {
-    const hexOnlyRegex = /^[0-9A-Fa-f]*$/;
-    // Check whether this is a trace ID or traceQL query by checking if it only contains hex characters
-    return !!query.trim().match(hexOnlyRegex);
-  }
-
   applyTemplateVariables(query: TempoQuery, scopedVars: ScopedVars) {
     return this.applyVariables(query, scopedVars);
   }
@@ -517,7 +517,104 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
     };
   }
 
-  handleMetricsSummary = (target: TempoQuery, query: string, options: DataQueryRequest<TempoQuery>) => {
+  formatGroupBy = (groupBy: TraceqlFilter[]) => {
+    return groupBy
+      ?.filter((f) => f.tag)
+      .map((f) => {
+        if (f.scope === TraceqlSearchScope.Unscoped) {
+          return `.${f.tag}`;
+        }
+        return f.scope !== TraceqlSearchScope.Intrinsic ? `${f.scope}.${f.tag}` : f.tag;
+      })
+      .join(', ');
+  };
+
+  hasGroupBy = (query: TempoQuery) => {
+    return query.groupBy?.find((gb) => gb.tag);
+  };
+
+  /**
+   * Handles the simplest of the queries where we have just a trace id and return trace data for it.
+   * @param options
+   * @param targets
+   * @private
+   */
+  handleTraceIdQuery(options: DataQueryRequest<TempoQuery>, targets: TempoQuery[]): Observable<DataQueryResponse> {
+    const validTargets = targets
+      .filter((t) => t.query)
+      .map((t): TempoQuery => ({ ...t, query: t.query?.trim(), queryType: 'traceId' }));
+    if (!validTargets.length) {
+      return EMPTY;
+    }
+
+    const request = this.makeTraceIdRequest(options, validTargets);
+    return super.query(request).pipe(
+      map((response) => {
+        if (response.error) {
+          return response;
+        }
+        return transformTrace(response, this.instanceSettings, this.nodeGraph?.enabled);
+      })
+    );
+  }
+
+  handleTraceQlQuery = (
+    options: DataQueryRequest<TempoQuery>,
+    targets: {
+      [type: string]: TempoQuery[];
+    },
+    queryValue: string
+  ): Observable<DataQueryResponse> => {
+    if (this.isStreamingSearchEnabled()) {
+      return this.handleStreamingQuery(options, targets.traceql, queryValue);
+    } else {
+      return this._request('/api/search', {
+        q: queryValue,
+        limit: options.targets[0].limit ?? DEFAULT_LIMIT,
+        spss: options.targets[0].spss ?? DEFAULT_SPSS,
+        start: options.range.from.unix(),
+        end: options.range.to.unix(),
+      }).pipe(
+        map((response) => {
+          return {
+            data: formatTraceQLResponse(response.data.traces, this.instanceSettings, targets.traceql[0].tableType),
+          };
+        }),
+        catchError((err) => {
+          return of({ error: { message: getErrorMessage(err.data.message) }, data: [] });
+        })
+      );
+    }
+  };
+
+  handleTraceQlMetricsQuery = (
+    options: DataQueryRequest<TempoQuery>,
+    queryValue: string
+  ): Observable<DataQueryResponse> => {
+    const requestData = {
+      query: queryValue,
+      start: options.range.from.unix(),
+      end: options.range.to.unix(),
+      step: options.targets[0].step,
+    };
+
+    if (!requestData.step) {
+      delete requestData.step;
+    }
+
+    return this._request('/api/metrics/query_range', requestData).pipe(
+      map((response) => {
+        return {
+          data: formatTraceQLMetrics(queryValue, response.data),
+        };
+      }),
+      catchError((err) => {
+        return of({ error: { message: getErrorMessage(err.data.message) }, data: [] });
+      })
+    );
+  };
+
+  handleMetricsSummaryQuery = (target: TempoQuery, query: string, options: DataQueryRequest<TempoQuery>) => {
     reportInteraction('grafana_traces_metrics_summary_queried', {
       datasourceType: 'tempo',
       app: options.app ?? '',
@@ -574,105 +671,30 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
     );
   };
 
-  formatGroupBy = (groupBy: TraceqlFilter[]) => {
-    return groupBy
-      ?.filter((f) => f.tag)
-      .map((f) => {
-        if (f.scope === TraceqlSearchScope.Unscoped) {
-          return `.${f.tag}`;
-        }
-        return f.scope !== TraceqlSearchScope.Intrinsic ? `${f.scope}.${f.tag}` : f.tag;
-      })
-      .join(', ');
-  };
-
-  hasGroupBy = (query: TempoQuery) => {
-    return query.groupBy?.find((gb) => gb.tag);
-  };
-
-  /**
-   * Handles the simplest of the queries where we have just a trace id and return trace data for it.
-   * @param options
-   * @param targets
-   * @private
-   */
-  handleTraceIdQuery(options: DataQueryRequest<TempoQuery>, targets: TempoQuery[]): Observable<DataQueryResponse> {
-    const validTargets = targets
-      .filter((t) => t.query)
-      .map((t): TempoQuery => ({ ...t, query: t.query?.trim(), queryType: 'traceId' }));
-    if (!validTargets.length) {
+  // This function can probably be simplified by avoiding passing both `targets` and `query`,
+  // since `query` is built from `targets`, if you look at how this function is currently called
+  handleStreamingQuery(
+    options: DataQueryRequest<TempoQuery>,
+    targets: TempoQuery[],
+    query: string
+  ): Observable<DataQueryResponse> {
+    if (query === '') {
       return EMPTY;
     }
 
-    const traceRequest = this.traceIdQueryRequest(options, validTargets);
-
-    return super.query(traceRequest).pipe(
-      map((response) => {
-        if (response.error) {
-          return response;
-        }
-        return transformTrace(response, this.instanceSettings, this.nodeGraph?.enabled);
-      })
+    return merge(
+      ...targets.map((target) =>
+        doTempoChannelStream(
+          { ...target, query },
+          this, // the datasource
+          options,
+          this.instanceSettings
+        )
+      )
     );
   }
 
-  handleTraceQlQuery = (
-    options: DataQueryRequest<TempoQuery>,
-    targets: {
-      [type: string]: TempoQuery[];
-    },
-    queryValue: string
-  ): Observable<DataQueryResponse> => {
-    if (this.isStreamingSearchEnabled()) {
-      return this.handleStreamingSearch(options, targets.traceql, queryValue);
-    } else {
-      return this._request('/api/search', {
-        q: queryValue,
-        limit: options.targets[0].limit ?? DEFAULT_LIMIT,
-        spss: options.targets[0].spss ?? DEFAULT_SPSS,
-        start: options.range.from.unix(),
-        end: options.range.to.unix(),
-      }).pipe(
-        map((response) => {
-          return {
-            data: formatTraceQLResponse(response.data.traces, this.instanceSettings, targets.traceql[0].tableType),
-          };
-        }),
-        catchError((err) => {
-          return of({ error: { message: getErrorMessage(err.data.message) }, data: [] });
-        })
-      );
-    }
-  };
-
-  handleTraceQlMetricsQuery = (
-    options: DataQueryRequest<TempoQuery>,
-    queryValue: string
-  ): Observable<DataQueryResponse> => {
-    const requestData = {
-      query: queryValue,
-      start: options.range.from.unix(),
-      end: options.range.to.unix(),
-      step: options.targets[0].step,
-    };
-
-    if (!requestData.step) {
-      delete requestData.step;
-    }
-
-    return this._request('/api/metrics/query_range', requestData).pipe(
-      map((response) => {
-        return {
-          data: formatTraceQLMetrics(queryValue, response.data),
-        };
-      }),
-      catchError((err) => {
-        return of({ error: { message: getErrorMessage(err.data.message) }, data: [] });
-      })
-    );
-  };
-
-  traceIdQueryRequest(options: DataQueryRequest<TempoQuery>, targets: TempoQuery[]): DataQueryRequest<TempoQuery> {
+  makeTraceIdRequest(options: DataQueryRequest<TempoQuery>, targets: TempoQuery[]): DataQueryRequest<TempoQuery> {
     const request = {
       ...options,
       targets,
@@ -697,29 +719,6 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
     return request;
   }
 
-  // This function can probably be simplified by avoiding passing both `targets` and `query`,
-  // since `query` is built from `targets`, if you look at how this function is currently called
-  handleStreamingSearch(
-    options: DataQueryRequest<TempoQuery>,
-    targets: TempoQuery[],
-    query: string
-  ): Observable<DataQueryResponse> {
-    if (query === '') {
-      return EMPTY;
-    }
-
-    return merge(
-      ...targets.map((target) =>
-        doTempoChannelStream(
-          { ...target, query },
-          this, // the datasource
-          options,
-          this.instanceSettings
-        )
-      )
-    );
-  }
-
   async metadataRequest(url: string, params = {}) {
     return await lastValueFrom(this._request(url, params, { method: 'GET', hideFromInspector: true }));
   }
@@ -728,7 +727,6 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
     const params = data ? urlUtil.serializeParams(data) : '';
     const url = `${this.instanceSettings.url}${apiUrl}${params.length ? `?${params}` : ''}`;
     const req = { ...options, url };
-
     return getBackendSrv().fetch(req);
   }
 
@@ -761,7 +759,7 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
       const from = new Date(now);
       from.setMinutes(from.getMinutes() - 15);
       observables.push(
-        this.handleStreamingSearch(
+        this.handleStreamingQuery(
           {
             range: {
               from: dateTime(from),
@@ -985,7 +983,7 @@ function errorAndDurationQuery(
         throw new Error(getErrorMessage(errorRes.error?.message));
       }
 
-      const serviceGraphView = getServiceGraphView(
+      const serviceGraphView = getServiceGraphViewDataFrames(
         request,
         rateResponse,
         errorAndDurationResponse[0],
@@ -1160,7 +1158,7 @@ function makePromServiceMapRequest(options: DataQueryRequest<TempoQuery>): DataQ
   };
 }
 
-function getServiceGraphView(
+function getServiceGraphViewDataFrames(
   request: DataQueryRequest<TempoQuery>,
   rateResponse: ServiceMapQueryResponseWithRates,
   secondResponse: DataQueryResponse,

public/app/plugins/datasource/tempo/graphTransform.test.ts

@@ -1,7 +1,7 @@
 import { DataFrameView, dateTime, createDataFrame, FieldType } from '@grafana/data';
 
 import { createGraphFrames, mapPromMetricsToServiceMap } from './graphTransform';
-import { bigResponse } from './testResponse';
+import { bigResponse } from './test/testResponse';
 
 describe('createGraphFrames', () => {
   it('transforms basic response into nodes and edges frame', async () => {

public/app/plugins/datasource/tempo/resultTransformer.test.ts

@@ -14,7 +14,7 @@ import {
   otlpDataFrameFromResponse,
   otlpResponse,
   traceQlResponse,
-} from './testResponse';
+} from './test/testResponse';
 import { TraceSearchMetadata } from './types';
 
 const defaultSettings: DataSourceInstanceSettings = {

public/app/plugins/datasource/tempo/test/mocks.ts

@@ -1,8 +1,8 @@
 import { DataSourceInstanceSettings, PluginType, toUtc } from '@grafana/data';
 import { TemplateSrv } from '@grafana/runtime';
 
-import { TempoDatasource } from './datasource';
-import { TempoJsonData } from './types';
+import { TempoDatasource } from '../datasource';
+import { TempoJsonData } from '../types';
 
 const rawRange = {
   from: toUtc('2018-04-25 10:00'),

public/app/plugins/datasource/tempo/variables.test.ts

@@ -3,7 +3,7 @@ import { lastValueFrom } from 'rxjs';
 import { DataQueryRequest, TimeRange } from '@grafana/data';
 
 import { TempoVariableQuery } from './VariableQueryEditor';
-import { createMetadataRequest, createTempoDatasource } from './mocks';
+import { createMetadataRequest, createTempoDatasource } from './test/mocks';
 import { TempoVariableSupport } from './variables';
 
 describe('TempoVariableSupport', () => {