QueryVariableEditor: Select a variable ds does not work (#83144)

public/app/features/dashboard-scene/settings/variables/components/QueryVariableForm.test.tsx

@@ -1,8 +1,7 @@
-import { render, screen, waitFor } from '@testing-library/react';
+import { act, render, screen, waitFor } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import React, { FormEvent } from 'react';
 import { of } from 'rxjs';
-import { MockDataSourceApi } from 'test/mocks/datasource_srv';
 
 import {
   LoadingState,
@@ -17,6 +16,7 @@ import { setRunRequest } from '@grafana/runtime';
 import { VariableRefresh, VariableSort } from '@grafana/schema';
 import { mockDataSource } from 'app/features/alerting/unified/mocks';
 import { LegacyVariableQueryEditor } from 'app/features/variables/editor/LegacyVariableQueryEditor';
+import { getVariableQueryEditor } from 'app/features/variables/editor/getVariableQueryEditor';
 
 import { QueryVariableEditorForm } from './QueryVariableForm';
 
@@ -42,7 +42,7 @@ jest.mock('@grafana/runtime/src/services/dataSourceSrv', () => ({
       },
     }),
     getList: () => [defaultDatasource, promDatasource],
-    getInstanceSettings: () => ({ ...defaultDatasource }),
+    getInstanceSettings: (uid: string) => (uid === promDatasource.uid ? promDatasource : defaultDatasource),
   }),
 }));
 
@@ -60,6 +60,11 @@ const runRequestMock = jest.fn().mockReturnValue(
 
 setRunRequest(runRequestMock);
 
+jest.mock('app/features/variables/editor/getVariableQueryEditor', () => ({
+  ...jest.requireActual('app/features/variables/editor/getVariableQueryEditor'),
+  getVariableQueryEditor: jest.fn(),
+}));
+
 describe('QueryVariableEditorForm', () => {
   const mockOnDataSourceChange = jest.fn();
   const mockOnQueryChange = jest.fn();
@@ -71,14 +76,13 @@ describe('QueryVariableEditorForm', () => {
   const mockOnIncludeAllChange = jest.fn();
   const mockOnAllValueChange = jest.fn();
 
-  const defaultProps = {
-    datasource: new MockDataSourceApi(promDatasource.name, undefined, promDatasource.meta),
+  const defaultProps: React.ComponentProps<typeof QueryVariableEditorForm> = {
+    datasource: { uid: defaultDatasource.uid, type: defaultDatasource.type },
     onDataSourceChange: mockOnDataSourceChange,
     query: 'my-query',
     onQueryChange: mockOnQueryChange,
     onLegacyQueryChange: mockOnLegacyQueryChange,
     timeRange: getDefaultTimeRange(),
-    VariableQueryEditor: LegacyVariableQueryEditor,
     regex: '.*',
     onRegExChange: mockOnRegExChange,
     sort: VariableSort.alphabeticalAsc,
@@ -93,9 +97,10 @@ describe('QueryVariableEditorForm', () => {
     onAllValueChange: mockOnAllValueChange,
   };
 
-  function setup(props?: React.ComponentProps<typeof QueryVariableEditorForm>) {
+  async function setup(props?: React.ComponentProps<typeof QueryVariableEditorForm>) {
+    jest.mocked(getVariableQueryEditor).mockResolvedValue(LegacyVariableQueryEditor);
     return {
-      renderer: render(<QueryVariableEditorForm {...defaultProps} {...props} />),
+      renderer: await act(() => render(<QueryVariableEditorForm {...defaultProps} {...props} />)),
       user: userEvent.setup(),
     };
   }
@@ -104,10 +109,10 @@ describe('QueryVariableEditorForm', () => {
     jest.clearAllMocks();
   });
 
-  it('should render the component with initializing the components correctly', () => {
+  it('should render the component with initializing the components correctly', async () => {
     const {
       renderer: { getByTestId, getByRole },
-    } = setup();
+    } = await setup();
     const dataSourcePicker = getByTestId(selectors.components.DataSourcePicker.inputV2);
     //const queryEditor = getByTestId('query-editor');
     const regexInput = getByTestId(
@@ -149,7 +154,7 @@ describe('QueryVariableEditorForm', () => {
   it('should call onDataSourceChange when changing the datasource', async () => {
     const {
       renderer: { getByTestId },
-    } = setup();
+    } = await setup();
     const dataSourcePicker = getByTestId(selectors.components.DataSourcePicker.inputV2);
     await userEvent.click(dataSourcePicker);
     await userEvent.click(screen.getByText(/prometheus/i));
@@ -161,7 +166,7 @@ describe('QueryVariableEditorForm', () => {
   it('should call onQueryChange when changing the query', async () => {
     const {
       renderer: { getByTestId },
-    } = setup();
+    } = await setup();
     const queryEditor = getByTestId(
       selectors.pages.Dashboard.Settings.Variables.Edit.QueryVariable.queryOptionsQueryInput
     );
@@ -178,7 +183,7 @@ describe('QueryVariableEditorForm', () => {
   it('should call onRegExChange when changing the regex', async () => {
     const {
       renderer: { getByTestId },
-    } = setup();
+    } = await setup();
     const regexInput = getByTestId(
       selectors.pages.Dashboard.Settings.Variables.Edit.QueryVariable.queryOptionsRegExInputV2
     );
@@ -193,7 +198,7 @@ describe('QueryVariableEditorForm', () => {
   it('should call onSortChange when changing the sort', async () => {
     const {
       renderer: { getByTestId },
-    } = setup();
+    } = await setup();
     const sortSelect = getByTestId(
       selectors.pages.Dashboard.Settings.Variables.Edit.QueryVariable.queryOptionsSortSelectV2
     );
@@ -211,7 +216,7 @@ describe('QueryVariableEditorForm', () => {
   it('should call onRefreshChange when changing the refresh', async () => {
     const {
       renderer: { getByTestId },
-    } = setup();
+    } = await setup();
     const refreshSelect = getByTestId(
       selectors.pages.Dashboard.Settings.Variables.Edit.QueryVariable.queryOptionsRefreshSelectV2
     );
@@ -226,7 +231,7 @@ describe('QueryVariableEditorForm', () => {
   it('should call onMultiChange when changing the multi switch', async () => {
     const {
       renderer: { getByTestId },
-    } = setup();
+    } = await setup();
     const multiSwitch = getByTestId(
       selectors.pages.Dashboard.Settings.Variables.Edit.General.selectionOptionsMultiSwitch
     );
@@ -240,7 +245,7 @@ describe('QueryVariableEditorForm', () => {
   it('should call onIncludeAllChange when changing the include all switch', async () => {
     const {
       renderer: { getByTestId },
-    } = setup();
+    } = await setup();
     const includeAllSwitch = getByTestId(
       selectors.pages.Dashboard.Settings.Variables.Edit.General.selectionOptionsIncludeAllSwitch
     );
@@ -254,7 +259,7 @@ describe('QueryVariableEditorForm', () => {
   it('should call onAllValueChange when changing the all value', async () => {
     const {
       renderer: { getByTestId },
-    } = setup();
+    } = await setup();
     const allValueInput = getByTestId(
       selectors.pages.Dashboard.Settings.Variables.Edit.General.selectionOptionsCustomAllInput
     );

public/app/features/dashboard-scene/settings/variables/components/QueryVariableForm.tsx

@@ -1,16 +1,18 @@
 import React, { FormEvent } from 'react';
+import { useAsync } from 'react-use';
 
-import { DataSourceApi, DataSourceInstanceSettings, SelectableValue, TimeRange } from '@grafana/data';
+import { DataSourceInstanceSettings, SelectableValue, TimeRange } from '@grafana/data';
 import { selectors } from '@grafana/e2e-selectors';
+import { getDataSourceSrv } from '@grafana/runtime';
 import { QueryVariable } from '@grafana/scenes';
-import { VariableRefresh, VariableSort } from '@grafana/schema';
+import { DataSourceRef, VariableRefresh, VariableSort } from '@grafana/schema';
 import { Field } from '@grafana/ui';
 import { QueryEditor } from 'app/features/dashboard-scene/settings/variables/components/QueryEditor';
 import { SelectionOptionsForm } from 'app/features/dashboard-scene/settings/variables/components/SelectionOptionsForm';
 import { DataSourcePicker } from 'app/features/datasources/components/picker/DataSourcePicker';
+import { getVariableQueryEditor } from 'app/features/variables/editor/getVariableQueryEditor';
 import { QueryVariableRefreshSelect } from 'app/features/variables/query/QueryVariableRefreshSelect';
 import { QueryVariableSortSelect } from 'app/features/variables/query/QueryVariableSortSelect';
-import { VariableQueryEditorType } from 'app/features/variables/types';
 
 import { VariableLegend } from './VariableLegend';
 import { VariableTextAreaField } from './VariableTextAreaField';
@@ -18,12 +20,11 @@ import { VariableTextAreaField } from './VariableTextAreaField';
 type VariableQueryType = QueryVariable['state']['query'];
 
 interface QueryVariableEditorFormProps {
-  datasource: DataSourceApi | undefined;
+  datasource?: DataSourceRef;
   onDataSourceChange: (dsSettings: DataSourceInstanceSettings) => void;
   query: VariableQueryType;
   onQueryChange: (query: VariableQueryType) => void;
   onLegacyQueryChange: (query: VariableQueryType, definition: string) => void;
-  VariableQueryEditor: VariableQueryEditorType | undefined;
   timeRange: TimeRange;
   regex: string | null;
   onRegExChange: (event: FormEvent<HTMLTextAreaElement>) => void;
@@ -40,12 +41,11 @@ interface QueryVariableEditorFormProps {
 }
 
 export function QueryVariableEditorForm({
-  datasource,
+  datasource: datasourceRef,
   onDataSourceChange,
   query,
   onQueryChange,
   onLegacyQueryChange,
-  VariableQueryEditor,
   timeRange,
   regex,
   onRegExChange,
@@ -60,11 +60,19 @@ export function QueryVariableEditorForm({
   allValue,
   onAllValueChange,
 }: QueryVariableEditorFormProps) {
+  const { value: dsConfig } = useAsync(async () => {
+    const datasource = await getDataSourceSrv().get(datasourceRef ?? '');
+    const VariableQueryEditor = await getVariableQueryEditor(datasource);
+
+    return { datasource, VariableQueryEditor };
+  }, [datasourceRef]);
+  const { datasource, VariableQueryEditor } = dsConfig ?? {};
+
   return (
     <>
       <VariableLegend>Query options</VariableLegend>
       <Field label="Data source" htmlFor="data-source-picker">
-        <DataSourcePicker current={datasource} onChange={onDataSourceChange} variables={true} width={30} />
+        <DataSourcePicker current={datasourceRef} onChange={onDataSourceChange} variables={true} width={30} />
       </Field>
 
       {datasource && VariableQueryEditor && (

public/app/features/dashboard-scene/settings/variables/editors/QueryVariableEditor.tsx

@@ -1,11 +1,8 @@
 import React, { FormEvent } from 'react';
-import { useAsync } from 'react-use';
 
 import { SelectableValue, DataSourceInstanceSettings } from '@grafana/data';
-import { getDataSourceSrv } from '@grafana/runtime';
 import { QueryVariable, sceneGraph } from '@grafana/scenes';
 import { DataSourceRef, VariableRefresh, VariableSort } from '@grafana/schema';
-import { getVariableQueryEditor } from 'app/features/variables/editor/getVariableQueryEditor';
 
 import { QueryVariableEditorForm } from '../components/QueryVariableForm';
 
@@ -16,17 +13,9 @@ interface QueryVariableEditorProps {
 type VariableQueryType = QueryVariable['state']['query'];
 
 export function QueryVariableEditor({ variable, onRunQuery }: QueryVariableEditorProps) {
-  const { datasource: datasourceRef, regex, sort, refresh, isMulti, includeAll, allValue, query } = variable.useState();
+  const { datasource, regex, sort, refresh, isMulti, includeAll, allValue, query } = variable.useState();
   const { value: timeRange } = sceneGraph.getTimeRange(variable).useState();
 
-  const { value: dsConfig } = useAsync(async () => {
-    const datasource = await getDataSourceSrv().get(datasourceRef ?? '');
-    const VariableQueryEditor = await getVariableQueryEditor(datasource);
-
-    return { datasource, VariableQueryEditor };
-  }, [datasourceRef]);
-  const { datasource, VariableQueryEditor } = dsConfig ?? {};
-
   const onRegExChange = (event: React.FormEvent<HTMLTextAreaElement>) => {
     variable.setState({ regex: event.currentTarget.value });
   };
@@ -56,12 +45,11 @@ export function QueryVariableEditor({ variable, onRunQuery }: QueryVariableEdito
 
   return (
     <QueryVariableEditorForm
-      datasource={datasource}
+      datasource={datasource ?? undefined}
       onDataSourceChange={onDataSourceChange}
       query={query}
       onQueryChange={onQueryChange}
       onLegacyQueryChange={onQueryChange}
-      VariableQueryEditor={VariableQueryEditor}
       timeRange={timeRange}
       regex={regex}
       onRegExChange={onRegExChange}

public/app/features/variables/query/QueryVariableEditor.test.tsx

@@ -1,8 +1,9 @@
-import { render, screen } from '@testing-library/react';
+import { render, screen, act } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import React from 'react';
+import { MockDataSourceApi } from 'test/mocks/datasource_srv';
 
-import { DataSourceApi, VariableSupportType } from '@grafana/data';
+import { VariableSupportType } from '@grafana/data';
 import { selectors } from '@grafana/e2e-selectors';
 import { mockDataSource } from 'app/features/alerting/unified/mocks';
 import { DataSourceType } from 'app/features/alerting/unified/utils/datasource';
@@ -15,11 +16,25 @@ import { QueryVariableModel } from '../types';
 import { Props, QueryVariableEditorUnConnected } from './QueryVariableEditor';
 import { initialQueryVariableModelState } from './reducer';
 
-const setupTestContext = (options: Partial<Props>) => {
+const mockDS = mockDataSource({
+  name: 'CloudManager',
+  type: DataSourceType.Alertmanager,
+});
+const ds = new MockDataSourceApi(mockDS);
+const editor = jest.fn().mockImplementation(LegacyVariableQueryEditor);
+
+ds.variables = {
+  getType: () => VariableSupportType.Custom,
+  query: jest.fn(),
+  editor: editor,
+  getDefaultQuery: jest.fn(),
+};
+
+const setupTestContext = async (options: Partial<Props>) => {
   const variableDefaults: Partial<QueryVariableModel> = { rootStateKey: 'key' };
   const extended = {
     VariableQueryEditor: LegacyVariableQueryEditor,
-    dataSource: {} as unknown as DataSourceApi,
+    dataSource: ds,
   };
 
   const defaults: Props = {
@@ -33,20 +48,15 @@ const setupTestContext = (options: Partial<Props>) => {
   };
 
   const props: Props & Record<string, unknown> = { ...defaults, ...options };
-  const { rerender } = render(<QueryVariableEditorUnConnected {...props} />);
+  const { rerender } = await act(() => render(<QueryVariableEditorUnConnected {...props} />));
 
   return { rerender, props };
 };
 
-const mockDS = mockDataSource({
-  name: 'CloudManager',
-  type: DataSourceType.Alertmanager,
-});
-
 jest.mock('@grafana/runtime/src/services/dataSourceSrv', () => {
   return {
     getDataSourceSrv: () => ({
-      get: () => Promise.resolve(mockDS),
+      get: async () => ds,
       getList: () => [mockDS],
       getInstanceSettings: () => mockDS,
     }),
@@ -57,8 +67,8 @@ const defaultIdentifier: KeyedVariableIdentifier = { type: 'query', rootStateKey
 
 describe('QueryVariableEditor', () => {
   describe('when the component is mounted', () => {
-    it('then it should call initQueryVariableEditor', () => {
-      const { props } = setupTestContext({});
+    it('then it should call initQueryVariableEditor', async () => {
+      const { props } = await setupTestContext({});
 
       expect(props.initQueryVariableEditor).toHaveBeenCalledTimes(1);
       expect(props.initQueryVariableEditor).toHaveBeenCalledWith(defaultIdentifier);
@@ -66,50 +76,29 @@ describe('QueryVariableEditor', () => {
   });
 
   describe('when the editor is rendered', () => {
-    const extendedCustom = {
-      extended: {
-        VariableQueryEditor: jest.fn().mockImplementation(LegacyVariableQueryEditor),
-        dataSource: {
-          variables: {
-            getType: () => VariableSupportType.Custom,
-            query: jest.fn(),
-            editor: jest.fn(),
-          },
-        } as unknown as DataSourceApi,
-      },
-    };
-    it('should pass down the query with default values if the datasource config defines it', () => {
-      const extended = { ...extendedCustom };
-      extended.extended.dataSource.variables!.getDefaultQuery = jest
-        .fn()
-        .mockImplementation(() => 'some default query');
-      const { props } = setupTestContext(extended);
-      expect(props.extended?.dataSource?.variables?.getDefaultQuery).toBeDefined();
-      expect(props.extended?.dataSource?.variables?.getDefaultQuery).toHaveBeenCalledTimes(1);
-      expect(props.extended?.VariableQueryEditor).toHaveBeenCalledWith(
-        expect.objectContaining({ query: 'some default query' }),
-        expect.anything()
-      );
+    beforeEach(() => {
+      jest.clearAllMocks();
     });
-    it('should not pass down a default query if the datasource config doesnt define it', () => {
-      extendedCustom.extended.dataSource.variables!.getDefaultQuery = undefined;
-      const { props } = setupTestContext(extendedCustom);
-      expect(props.extended?.dataSource?.variables?.getDefaultQuery).not.toBeDefined();
-      expect(props.extended?.VariableQueryEditor).toHaveBeenCalledWith(
-        expect.objectContaining({ query: '' }),
-        expect.anything()
-      );
+
+    it('should pass down the query with default values if the datasource config defines it', async () => {
+      ds.variables!.getDefaultQuery = jest.fn().mockImplementationOnce(() => 'some default query');
+
+      await setupTestContext({});
+      expect(ds.variables?.getDefaultQuery).toBeDefined();
+      expect(ds.variables?.getDefaultQuery).toHaveBeenCalledTimes(1);
+      expect(editor.mock.calls[0][0].query).toBe('some default query');
     });
   });
+
   describe('when the user changes', () => {
     it.each`
       fieldName  | propName                      | expectedArgs
-      ${'query'} | ${'changeQueryVariableQuery'} | ${[defaultIdentifier, 't', 't']}
+      ${'query'} | ${'changeQueryVariableQuery'} | ${[defaultIdentifier, 't', '']}
       ${'regex'} | ${'onPropChange'}             | ${[{ propName: 'regex', propValue: 't', updateOptions: true }]}
     `(
       '$fieldName field and tabs away then $propName should be called with correct args',
       async ({ fieldName, propName, expectedArgs }) => {
-        const { props } = setupTestContext({});
+        const { props } = await setupTestContext({});
         const propUnderTest = props[propName];
         const fieldAccessor = fieldAccessors[fieldName];
 
@@ -130,7 +119,7 @@ describe('QueryVariableEditor', () => {
     `(
       '$fieldName field but reverts the change and tabs away then $propName should not be called',
       async ({ fieldName, propName }) => {
-        const { props } = setupTestContext({});
+        const { props } = await setupTestContext({});
         const propUnderTest = props[propName];
         const fieldAccessor = fieldAccessors[fieldName];
 

public/app/features/variables/query/QueryVariableEditor.tsx

@@ -125,23 +125,19 @@ export class QueryVariableEditorUnConnected extends PureComponent<Props, State>
 
   render() {
     const { extended, variable } = this.props;
-
     if (!extended || !extended.dataSource || !extended.VariableQueryEditor) {
       return null;
     }
 
-    const datasource = extended.dataSource;
-    const VariableQueryEditor = extended.VariableQueryEditor;
     const timeRange = getTimeSrv().timeRange();
 
     return (
       <QueryVariableEditorForm
-        datasource={datasource}
+        datasource={variable.datasource ?? undefined}
         onDataSourceChange={this.onDataSourceChange}
         query={variable.query}
         onQueryChange={this.onQueryChange}
         onLegacyQueryChange={this.onLegacyQueryChange}
-        VariableQueryEditor={VariableQueryEditor}
         timeRange={timeRange}
         regex={variable.regex}
         onRegExChange={this.onRegExBlur}

public/test/mocks/datasource_srv.ts

@@ -36,12 +36,14 @@ export class MockDataSourceApi extends DataSourceApi {
   result: DataQueryResponse = { data: [] };
 
   constructor(
-    name?: string,
+    datasource: string | DataSourceInstanceSettings = 'MockDataSourceApi',
     result?: DataQueryResponse,
     meta?: DataSourcePluginMeta,
     public error: string | null = null
   ) {
-    super({ name: name ? name : 'MockDataSourceApi' } as DataSourceInstanceSettings);
+    const instaceSettings =
+      typeof datasource === 'string' ? ({ name: datasource } as DataSourceInstanceSettings) : datasource;
+    super(instaceSettings);
     if (result) {
       this.result = result;
     }