Azure Monitor: Fix subscription selector when changing data sources (#…

…56284) (#56408)

(cherry picked from commit 2cdbd05)

Co-authored-by: Andres Martinez Gotor <andres.martinez@grafana.com>

public/app/plugins/datasource/grafana-azure-monitor-datasource/components/ArgQueryEditor/ArgQueryEditor.test.tsx

@@ -0,0 +1,84 @@
+import { render, screen } from '@testing-library/react';
+import React from 'react';
+
+import createMockDatasource from '../../__mocks__/datasource';
+import createMockQuery from '../../__mocks__/query';
+
+import ArgQueryEditor from './ArgQueryEditor';
+
+jest.mock('@grafana/runtime', () => ({
+  ...(jest.requireActual('@grafana/runtime') as unknown as object),
+  getTemplateSrv: () => ({
+    replace: (val: string) => {
+      return val;
+    },
+  }),
+}));
+
+const variableOptionGroup = {
+  label: 'Template variables',
+  options: [],
+};
+
+const defaultProps = {
+  query: createMockQuery(),
+  datasource: createMockDatasource(),
+  variableOptionGroup: variableOptionGroup,
+  onChange: jest.fn(),
+  setError: jest.fn(),
+};
+
+describe('ArgQueryEditor', () => {
+  it('should render', async () => {
+    render(<ArgQueryEditor {...defaultProps} />);
+    expect(await screen.findByTestId('azure-monitor-arg-query-editor-with-experimental-ui')).toBeInTheDocument();
+  });
+
+  it('should select a subscription from the fetched array', async () => {
+    const datasource = createMockDatasource({
+      getSubscriptions: jest.fn().mockResolvedValue([{ value: 'foo' }]),
+    });
+    const onChange = jest.fn();
+    render(<ArgQueryEditor {...defaultProps} datasource={datasource} onChange={onChange} />);
+    expect(await screen.findByTestId('azure-monitor-arg-query-editor-with-experimental-ui')).toBeInTheDocument();
+    expect(onChange).toHaveBeenCalledWith(expect.objectContaining({ subscriptions: ['foo'] }));
+  });
+
+  it('should select a subscription from existing query', async () => {
+    const onChange = jest.fn();
+    const query = createMockQuery({
+      subscriptions: ['bar'],
+    });
+    render(<ArgQueryEditor {...defaultProps} onChange={onChange} query={query} />);
+    expect(await screen.findByTestId('azure-monitor-arg-query-editor-with-experimental-ui')).toBeInTheDocument();
+    expect(onChange).toHaveBeenCalledWith(expect.objectContaining({ subscriptions: ['bar'] }));
+  });
+
+  it('should change the subscription if the selected one is not part of the fetched array', async () => {
+    const onChange = jest.fn();
+    const datasource = createMockDatasource({
+      getSubscriptions: jest.fn().mockResolvedValue([{ value: 'foo' }]),
+    });
+    const query = createMockQuery({
+      subscriptions: ['bar'],
+    });
+    render(<ArgQueryEditor {...defaultProps} datasource={datasource} onChange={onChange} query={query} />);
+    expect(await screen.findByTestId('azure-monitor-arg-query-editor-with-experimental-ui')).toBeInTheDocument();
+    expect(onChange).toHaveBeenCalledWith(expect.objectContaining({ subscriptions: ['foo'] }));
+    expect(onChange).not.toHaveBeenCalledWith(expect.objectContaining({ subscriptions: ['bar'] }));
+  });
+
+  it('should keep a subset of subscriptions if the new list does not contain all the query subscriptions', async () => {
+    const onChange = jest.fn();
+    const datasource = createMockDatasource({
+      getSubscriptions: jest.fn().mockResolvedValue([{ value: 'foo' }, { value: 'bar' }]),
+    });
+    const query = createMockQuery({
+      subscriptions: ['foo', 'bar', 'foobar'],
+    });
+    render(<ArgQueryEditor {...defaultProps} datasource={datasource} onChange={onChange} query={query} />);
+    expect(await screen.findByTestId('azure-monitor-arg-query-editor-with-experimental-ui')).toBeInTheDocument();
+    expect(onChange).toHaveBeenCalledWith(expect.objectContaining({ subscriptions: ['foo', 'bar'] }));
+    expect(onChange).not.toHaveBeenCalledWith(expect.objectContaining({ subscriptions: ['foo', 'bar', 'foobar'] }));
+  });
+});

public/app/plugins/datasource/grafana-azure-monitor-datasource/components/ArgQueryEditor/ArgQueryEditor.tsx

@@ -1,4 +1,5 @@
-import React, { useEffect, useState, useRef } from 'react';
+import { intersection } from 'lodash';
+import React, { useState, useMemo } from 'react';
 
 import { EditorRows, EditorRow, EditorFieldGroup } from '@grafana/experimental';
 
@@ -18,6 +19,28 @@ interface ArgQueryEditorProps {
 }
 
 const ERROR_SOURCE = 'arg-subscriptions';
+
+function selectSubscriptions(
+  fetchedSubscriptions: string[],
+  currentSubscriptions?: string[],
+  currentSubscription?: string
+) {
+  let querySubscriptions = currentSubscriptions || [];
+  if (querySubscriptions.length === 0 && currentSubscription) {
+    querySubscriptions = [currentSubscription];
+  }
+  if (querySubscriptions.length === 0 && fetchedSubscriptions.length) {
+    querySubscriptions = [fetchedSubscriptions[0]];
+  }
+  const commonSubscriptions = intersection(querySubscriptions, fetchedSubscriptions);
+  if (fetchedSubscriptions.length && querySubscriptions.length > commonSubscriptions.length) {
+    // If not all of the query subscriptions are in the list of fetched subscriptions, then
+    // select only the ones present (or the first one if none is present)
+    querySubscriptions = commonSubscriptions.length > 0 ? commonSubscriptions : [fetchedSubscriptions[0]];
+  }
+  return querySubscriptions;
+}
+
 const ArgQueryEditor: React.FC<ArgQueryEditorProps> = ({
   query,
   datasource,
@@ -26,31 +49,28 @@ const ArgQueryEditor: React.FC<ArgQueryEditorProps> = ({
   onChange,
   setError,
 }) => {
-  const fetchedRef = useRef(false);
   const [subscriptions, setSubscriptions] = useState<AzureMonitorOption[]>([]);
-
-  useEffect(() => {
-    if (fetchedRef.current) {
-      return;
-    }
-
-    fetchedRef.current = true;
+  useMemo(() => {
     datasource
       .getSubscriptions()
       .then((results) => {
         const fetchedSubscriptions = results.map((v) => ({ label: v.text, value: v.value, description: v.value }));
         setSubscriptions(fetchedSubscriptions);
         setError(ERROR_SOURCE, undefined);
 
-        if (!query.subscriptions?.length && fetchedSubscriptions?.length) {
-          onChange({
-            ...query,
-            subscriptions: [query.subscription ?? fetchedSubscriptions[0].value],
-          });
-        }
+        onChange({
+          ...query,
+          subscriptions: selectSubscriptions(
+            fetchedSubscriptions.map((v) => v.value),
+            query.subscriptions,
+            query.subscription
+          ),
+        });
       })
       .catch((err) => setError(ERROR_SOURCE, err));
-  }, [datasource, onChange, query, setError]);
+    // We are only interested in re-fetching subscriptions if the data source changes
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [datasource]);
 
   return (
     <span data-testid="azure-monitor-arg-query-editor-with-experimental-ui">