From 038c97f31c4adf8c58263918d71c8e69a34745e4 Mon Sep 17 00:00:00 2001
From: Andrej Ocenas <mr.ocenas@gmail.com>
Date: Tue, 29 Nov 2022 22:50:01 +0100
Subject: [PATCH] TestDatasource: Add scenario for generating trace data
 (#59299)

---
 pkg/tsdb/testdatasource/scenarios.go          |  6 +++
 .../datasource/testdata/QueryEditor.tsx       | 12 ++++++
 .../plugins/datasource/testdata/datasource.ts | 42 +++++++++++++++++++
 .../app/plugins/datasource/testdata/types.ts  |  1 +
 4 files changed, 61 insertions(+)

diff --git a/pkg/tsdb/testdatasource/scenarios.go b/pkg/tsdb/testdatasource/scenarios.go
index adca0fd977e..07349c9bf84 100644
--- a/pkg/tsdb/testdatasource/scenarios.go
+++ b/pkg/tsdb/testdatasource/scenarios.go
@@ -43,6 +43,7 @@ const (
 	rawFrameQuery                     queryType = "raw_frame"
 	csvFileQueryType                  queryType = "csv_file"
 	csvContentQueryType               queryType = "csv_content"
+	traceType                         queryType = "trace"
 )
 
 type queryType string
@@ -218,6 +219,11 @@ Timestamps will line up evenly on timeStepSeconds (For example, 60 seconds means
 		handler: s.handleCsvContentScenario,
 	})
 
+	s.registerScenario(&Scenario{
+		ID:   string(traceType),
+		Name: "Trace",
+	})
+
 	s.queryMux.HandleFunc("", s.handleFallbackScenario)
 }
 
diff --git a/public/app/plugins/datasource/testdata/QueryEditor.tsx b/public/app/plugins/datasource/testdata/QueryEditor.tsx
index 084a6e0d125..b3c7d996b97 100644
--- a/public/app/plugins/datasource/testdata/QueryEditor.tsx
+++ b/public/app/plugins/datasource/testdata/QueryEditor.tsx
@@ -302,6 +302,18 @@ export const QueryEditor = ({ query, datasource, onChange, onRunQuery }: Props)
         <NodeGraphEditor onChange={(val: NodesQuery) => onChange({ ...query, nodes: val })} query={query} />
       )}
       {scenarioId === 'server_error_500' && <ErrorEditor onChange={onUpdate} query={query} ds={datasource} />}
+      {scenarioId === 'trace' && (
+        <InlineField labelWidth={14} label="Span count">
+          <Input
+            type="number"
+            name="spanCount"
+            value={query.spanCount}
+            width={32}
+            onChange={onInputChange}
+            placeholder="10"
+          />
+        </InlineField>
+      )}
 
       {description && <p>{description}</p>}
     </>
diff --git a/public/app/plugins/datasource/testdata/datasource.ts b/public/app/plugins/datasource/testdata/datasource.ts
index fd9e42ea665..de176ee6de5 100644
--- a/public/app/plugins/datasource/testdata/datasource.ts
+++ b/public/app/plugins/datasource/testdata/datasource.ts
@@ -14,6 +14,7 @@ import {
   TimeRange,
   ScopedVars,
   toDataFrame,
+  MutableDataFrame,
 } from '@grafana/data';
 import { DataSourceWithBackend, getBackendSrv, getGrafanaLiveSrv, getTemplateSrv, TemplateSrv } from '@grafana/runtime';
 import { getSearchFilterScopedVar } from 'app/features/variables/utils';
@@ -70,6 +71,9 @@ export class TestDataDataSource extends DataSourceWithBackend<TestDataQuery> {
         case 'flame_graph':
           streams.push(this.flameGraphQuery());
           break;
+        case 'trace':
+          streams.push(this.trace(target, options));
+          break;
         case 'raw_frame':
           streams.push(this.rawFrameQuery(target, options));
           break;
@@ -224,6 +228,44 @@ export class TestDataDataSource extends DataSourceWithBackend<TestDataQuery> {
     return of({ data: [flameGraphData] }).pipe(delay(100));
   }
 
+  trace(target: TestDataQuery, options: DataQueryRequest<TestDataQuery>): Observable<DataQueryResponse> {
+    const frame = new MutableDataFrame({
+      meta: {
+        preferredVisualisationType: 'trace',
+      },
+      fields: [
+        { name: 'traceID' },
+        { name: 'spanID' },
+        { name: 'parentSpanID' },
+        { name: 'operationName' },
+        { name: 'serviceName' },
+        { name: 'serviceTags' },
+        { name: 'startTime' },
+        { name: 'duration' },
+        { name: 'logs' },
+        { name: 'references' },
+        { name: 'tags' },
+      ],
+    });
+    const numberOfSpans = options.targets[0].spanCount || 10;
+    const spanIdPrefix = '75c665dfb68';
+    const start = Date.now() - 1000 * 60 * 30;
+
+    for (let i = 0; i < numberOfSpans; i++) {
+      frame.add({
+        traceID: spanIdPrefix + '10000',
+        spanID: spanIdPrefix + (10000 + i),
+        parentSpanID: i === 0 ? '' : spanIdPrefix + 10000,
+        operationName: `Operation ${i}`,
+        serviceName: `Service ${i}`,
+        startTime: start + i * 100,
+        duration: 300,
+      });
+    }
+
+    return of({ data: [frame] }).pipe(delay(100));
+  }
+
   rawFrameQuery(target: TestDataQuery, options: DataQueryRequest<TestDataQuery>): Observable<DataQueryResponse> {
     try {
       const data = JSON.parse(target.rawFrameContent ?? '[]').map((v: any) => {
diff --git a/public/app/plugins/datasource/testdata/types.ts b/public/app/plugins/datasource/testdata/types.ts
index 82495788dcc..da8328ba32b 100644
--- a/public/app/plugins/datasource/testdata/types.ts
+++ b/public/app/plugins/datasource/testdata/types.ts
@@ -27,6 +27,7 @@ export interface TestDataQuery extends DataQuery {
   seriesCount?: number;
   usa?: USAQuery;
   errorType?: 'server_panic' | 'frontend_exception' | 'frontend_observable';
+  spanCount?: number;
 }
 
 export interface NodesQuery {