grafana/public/app/features/scenes/querying/SceneQueryRunner.test.ts

import { map, of } from 'rxjs';

import {
  ArrayVector,
  DataQueryRequest,
  DataSourceApi,
  getDefaultTimeRange,
  LoadingState,
  PanelData,
  standardTransformersRegistry,
  toDataFrame,
} from '@grafana/data';

import { SceneTimeRange } from '../core/SceneTimeRange';

import { SceneQueryRunner } from './SceneQueryRunner';

const getDatasource = () => {
  return {
    getRef: () => ({ uid: 'test' }),
  };
};

jest.mock('app/features/plugins/datasource_srv', () => ({
  getDatasourceSrv: jest.fn(() => ({
    get: getDatasource,
  })),
}));

const runRequest = jest.fn().mockReturnValue(
  of<PanelData>({
    state: LoadingState.Done,
    series: [
      toDataFrame([
        [100, 1],
        [200, 2],
        [300, 3],
      ]),
    ],
    timeRange: getDefaultTimeRange(),
  })
);

let sentRequest: DataQueryRequest | undefined;

jest.mock('app/features/query/state/runRequest', () => ({
  runRequest: (ds: DataSourceApi, request: DataQueryRequest) => {
    sentRequest = request;
    return runRequest(ds, request);
  },
}));

describe('SceneQueryRunner', () => {
  describe('when activated and got no data', () => {
    it('should run queries', async () => {
      const queryRunner = new SceneQueryRunner({
        queries: [{ refId: 'A' }],
        $timeRange: new SceneTimeRange(),
      });

      expect(queryRunner.state.data).toBeUndefined();

      queryRunner.activate();

      await new Promise((r) => setTimeout(r, 1));

      expect(queryRunner.state.data?.state).toBe(LoadingState.Done);
      // Default max data points
      expect(sentRequest?.maxDataPoints).toBe(500);
    });
  });

  describe('when activated and maxDataPointsFromWidth set to true', () => {
    it('should run queries', async () => {
      const queryRunner = new SceneQueryRunner({
        queries: [{ refId: 'A' }],
        $timeRange: new SceneTimeRange(),
        maxDataPointsFromWidth: true,
      });

      expect(queryRunner.state.data).toBeUndefined();

      queryRunner.activate();

      await new Promise((r) => setTimeout(r, 1));

      expect(queryRunner.state.data?.state).toBeUndefined();

      queryRunner.setContainerWidth(1000);

      expect(queryRunner.state.data?.state).toBeUndefined();

      await new Promise((r) => setTimeout(r, 1));

      expect(queryRunner.state.data?.state).toBe(LoadingState.Done);
    });
  });

  describe('transformations', () => {
    let transformerSpy1 = jest.fn();
    let transformerSpy2 = jest.fn();

    beforeEach(() => {
      standardTransformersRegistry.setInit(() => {
        return [
          {
            id: 'customTransformer1',
            editor: () => null,
            transformation: {
              id: 'customTransformer1',
              name: 'Custom Transformer',
              operator: (options) => (source) => {
                transformerSpy1(options);
                return source.pipe(
                  map((data) => {
                    return data.map((frame) => {
                      return {
                        ...frame,
                        fields: frame.fields.map((field) => {
                          return {
                            ...field,
                            values: new ArrayVector(field.values.toArray().map((v) => v * 2)),
                          };
                        }),
                      };
                    });
                  })
                );
              },
            },
            name: 'Custom Transformer',
          },
          {
            id: 'customTransformer2',
            editor: () => null,
            transformation: {
              id: 'customTransformer2',
              name: 'Custom Transformer2',
              operator: (options) => (source) => {
                transformerSpy2(options);
                return source.pipe(
                  map((data) => {
                    return data.map((frame) => {
                      return {
                        ...frame,
                        fields: frame.fields.map((field) => {
                          return {
                            ...field,
                            values: new ArrayVector(field.values.toArray().map((v) => v * 3)),
                          };
                        }),
                      };
                    });
                  })
                );
              },
            },
            name: 'Custom Transformer 2',
          },
        ];
      });
    });

    it('should apply transformations to query results', async () => {
      const queryRunner = new SceneQueryRunner({
        queries: [{ refId: 'A' }],
        $timeRange: new SceneTimeRange(),
        maxDataPoints: 100,
        transformations: [
          {
            id: 'customTransformer1',
            options: {
              option: 'value1',
            },
          },
          {
            id: 'customTransformer2',
            options: {
              option: 'value2',
            },
          },
        ],
      });

      queryRunner.activate();

      await new Promise((r) => setTimeout(r, 1));

      expect(queryRunner.state.data?.state).toBe(LoadingState.Done);
      expect(transformerSpy1).toHaveBeenCalledTimes(1);
      expect(transformerSpy1).toHaveBeenCalledWith({ option: 'value1' });
      expect(transformerSpy2).toHaveBeenCalledTimes(1);
      expect(transformerSpy2).toHaveBeenCalledWith({ option: 'value2' });
      expect(queryRunner.state.data?.series).toHaveLength(1);
      expect(queryRunner.state.data?.series[0].fields).toHaveLength(2);
      expect(queryRunner.state.data?.series[0].fields[0].values.toArray()).toEqual([600, 1200, 1800]);
      expect(queryRunner.state.data?.series[0].fields[1].values.toArray()).toEqual([6, 12, 18]);
    });
  });
});