Azure Monitor: refactor resource picker (#48312)

* Refactor Resource Picker to split up data logic from view logic.

public/app/plugins/datasource/grafana-azure-monitor-datasource/__mocks__/resourcePickerData.ts

@@ -1,20 +0,0 @@
-import ResourcePicker from '../resourcePicker/resourcePickerData';
-
-type DeepPartial<T> = {
-  [P in keyof T]?: DeepPartial<T[P]>;
-};
-
-export default function createMockResourcePickerData(overrides?: DeepPartial<ResourcePicker>) {
-  const _mockResourcePicker: DeepPartial<ResourcePicker> = {
-    getSubscriptions: () => jest.fn().mockResolvedValue([]),
-    getResourceGroupsBySubscriptionId: jest.fn().mockResolvedValue([]),
-    getResourcesForResourceGroup: jest.fn().mockResolvedValue([]),
-    getResourceURIFromWorkspace: jest.fn().mockReturnValue(''),
-    getResourceURIDisplayProperties: jest.fn().mockResolvedValue({}),
-    ...overrides,
-  };
-
-  const mockDatasource = _mockResourcePicker as ResourcePicker;
-
-  return jest.mocked(mockDatasource, true);
-}

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

@@ -5,13 +5,14 @@ import React from 'react';
 import { selectOptionInTest } from '@grafana/ui';
 
 import createMockDatasource from '../../__mocks__/datasource';
+import { createMockInstanceSetttings } from '../../__mocks__/instanceSettings';
 import createMockQuery from '../../__mocks__/query';
-import createMockResourcePickerData from '../../__mocks__/resourcePickerData';
 import {
   createMockResourceGroupsBySubscription,
   createMockSubscriptions,
   mockResourcesByResourceGroup,
 } from '../../__mocks__/resourcePickerRows';
+import ResourcePickerData from '../../resourcePicker/resourcePickerData';
 
 import MetricsQueryEditor from './MetricsQueryEditor';
 
@@ -20,11 +21,19 @@ const variableOptionGroup = {
   options: [],
 };
 
-const resourcePickerData = createMockResourcePickerData({
-  getSubscriptions: jest.fn().mockResolvedValue(createMockSubscriptions()),
-  getResourceGroupsBySubscriptionId: jest.fn().mockResolvedValue(createMockResourceGroupsBySubscription()),
-  getResourcesForResourceGroup: jest.fn().mockResolvedValue(mockResourcesByResourceGroup()),
-});
+export function createMockResourcePickerData() {
+  const mockDatasource = new ResourcePickerData(createMockInstanceSetttings());
+
+  mockDatasource.getSubscriptions = jest.fn().mockResolvedValue(createMockSubscriptions());
+  mockDatasource.getResourceGroupsBySubscriptionId = jest
+    .fn()
+    .mockResolvedValue(createMockResourceGroupsBySubscription());
+  mockDatasource.getResourcesForResourceGroup = jest.fn().mockResolvedValue(mockResourcesByResourceGroup());
+  mockDatasource.getResourceURIFromWorkspace = jest.fn().mockReturnValue('');
+  mockDatasource.getResourceURIDisplayProperties = jest.fn().mockResolvedValue({});
+
+  return mockDatasource;
+}
 
 describe('MetricsQueryEditor', () => {
   const originalScrollIntoView = window.HTMLElement.prototype.scrollIntoView;
@@ -36,7 +45,7 @@ describe('MetricsQueryEditor', () => {
   });
 
   it('should render', async () => {
-    const mockDatasource = createMockDatasource({ resourcePickerData });
+    const mockDatasource = createMockDatasource({ resourcePickerData: createMockResourcePickerData() });
 
     render(
       <MetricsQueryEditor
@@ -52,7 +61,7 @@ describe('MetricsQueryEditor', () => {
   });
 
   it('should change resource when a resource is selected in the ResourcePicker', async () => {
-    const mockDatasource = createMockDatasource({ resourcePickerData });
+    const mockDatasource = createMockDatasource({ resourcePickerData: createMockResourcePickerData() });
     const query = createMockQuery();
     delete query?.azureMonitor?.resourceUri;
     const onChange = jest.fn();
@@ -101,7 +110,7 @@ describe('MetricsQueryEditor', () => {
   });
 
   it('should reset metric namespace, metric name, and aggregation fields after selecting a new resource when a valid query has already been set', async () => {
-    const mockDatasource = createMockDatasource({ resourcePickerData });
+    const mockDatasource = createMockDatasource({ resourcePickerData: createMockResourcePickerData() });
     const query = createMockQuery();
     const onChange = jest.fn();
 
@@ -159,7 +168,7 @@ describe('MetricsQueryEditor', () => {
   });
 
   it('should change the metric name when selected', async () => {
-    const mockDatasource = createMockDatasource({ resourcePickerData });
+    const mockDatasource = createMockDatasource({ resourcePickerData: createMockResourcePickerData() });
     const onChange = jest.fn();
     const mockQuery = createMockQuery();
     mockDatasource.azureMonitorDatasource.getMetricNames = jest.fn().mockResolvedValue([
@@ -199,7 +208,7 @@ describe('MetricsQueryEditor', () => {
   });
 
   it('should change the aggregation type when selected', async () => {
-    const mockDatasource = createMockDatasource({ resourcePickerData });
+    const mockDatasource = createMockDatasource({ resourcePickerData: createMockResourcePickerData() });
     const onChange = jest.fn();
     const mockQuery = createMockQuery();
 

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

@@ -2,12 +2,13 @@ import { act, render, screen } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import React from 'react';
 
-import createMockResourcePickerData from '../../__mocks__/resourcePickerData';
+import { createMockInstanceSetttings } from '../../__mocks__/instanceSettings';
 import {
   createMockResourceGroupsBySubscription,
   createMockSubscriptions,
   mockResourcesByResourceGroup,
 } from '../../__mocks__/resourcePickerRows';
+import ResourcePickerData from '../../resourcePicker/resourcePickerData';
 
 import { ResourceRowType } from './types';
 
@@ -20,14 +21,24 @@ const singleResourceSelectionURI =
   '/subscriptions/def-456/resourceGroups/dev-3/providers/Microsoft.Compute/virtualMachines/db-server';
 
 const noop: any = () => {};
+function createMockResourcePickerData() {
+  const mockDatasource = new ResourcePickerData(createMockInstanceSetttings());
+
+  mockDatasource.getSubscriptions = jest.fn().mockResolvedValue(createMockSubscriptions());
+  mockDatasource.getResourceGroupsBySubscriptionId = jest
+    .fn()
+    .mockResolvedValue(createMockResourceGroupsBySubscription());
+  mockDatasource.getResourcesForResourceGroup = jest.fn().mockResolvedValue(mockResourcesByResourceGroup());
+  mockDatasource.getResourceURIFromWorkspace = jest.fn().mockReturnValue('');
+  mockDatasource.getResourceURIDisplayProperties = jest.fn().mockResolvedValue({});
+
+  return mockDatasource;
+}
+
 const defaultProps = {
   templateVariables: [],
   resourceURI: noResourceURI,
-  resourcePickerData: createMockResourcePickerData({
-    getSubscriptions: jest.fn().mockResolvedValue(createMockSubscriptions()),
-    getResourceGroupsBySubscriptionId: jest.fn().mockResolvedValue(createMockResourceGroupsBySubscription()),
-    getResourcesForResourceGroup: jest.fn().mockResolvedValue(mockResourcesByResourceGroup()),
-  }),
+  resourcePickerData: createMockResourcePickerData(),
   onCancel: noop,
   onApply: noop,
   selectableEntryTypes: [

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

@@ -10,7 +10,7 @@ import { Space } from '../Space';
 import NestedRow from './NestedRow';
 import getStyles from './styles';
 import { ResourceRow, ResourceRowGroup, ResourceRowType } from './types';
-import { addResources, findRow, parseResourceURI } from './utils';
+import { findRow } from './utils';
 
 interface ResourcePickerProps {
   resourcePickerData: ResourcePickerData;
@@ -32,7 +32,7 @@ const ResourcePicker = ({
 
   type LoadingStatus = 'NotStarted' | 'Started' | 'Done';
   const [loadingStatus, setLoadingStatus] = useState<LoadingStatus>('NotStarted');
-  const [azureRows, setAzureRows] = useState<ResourceRowGroup>([]);
+  const [rows, setRows] = useState<ResourceRowGroup>([]);
   const [internalSelectedURI, setInternalSelectedURI] = useState<string | undefined>(resourceURI);
   const [errorMessage, setErrorMessage] = useState<string | undefined>(undefined);
   const [isAdvancedOpen, setIsAdvancedOpen] = useState(resourceURI?.includes('$'));
@@ -47,32 +47,8 @@ const ResourcePicker = ({
       const loadInitialData = async () => {
         try {
           setLoadingStatus('Started');
-          let resources = await resourcePickerData.getSubscriptions();
-          if (!internalSelectedURI) {
-            setAzureRows(resources);
-            setLoadingStatus('Done');
-            return;
-          }
-
-          const parsedURI = parseResourceURI(internalSelectedURI ?? '');
-          if (parsedURI) {
-            const resourceGroupURI = `/subscriptions/${parsedURI.subscriptionID}/resourceGroups/${parsedURI.resourceGroup}`;
-
-            // if a resource group was previously selected, but the resource groups under the parent subscription have not been loaded yet
-            if (parsedURI.resourceGroup && !findRow(resources, resourceGroupURI)) {
-              const resourceGroups = await resourcePickerData.getResourceGroupsBySubscriptionId(
-                parsedURI.subscriptionID
-              );
-              resources = addResources(resources, `/subscriptions/${parsedURI.subscriptionID}`, resourceGroups);
-            }
-
-            // if a resource was previously selected, but the resources under the parent resource group have not been loaded yet
-            if (parsedURI.resource && !findRow(azureRows, parsedURI.resource ?? '')) {
-              const resourcesForResourceGroup = await resourcePickerData.getResourcesForResourceGroup(resourceGroupURI);
-              resources = addResources(resources, resourceGroupURI, resourcesForResourceGroup);
-            }
-          }
-          setAzureRows(resources);
+          const resources = await resourcePickerData.fetchInitialRows(internalSelectedURI || '');
+          setRows(resources);
           setLoadingStatus('Done');
         } catch (error) {
           setLoadingStatus('Done');
@@ -82,11 +58,11 @@ const ResourcePicker = ({
 
       loadInitialData();
     }
-  }, [resourcePickerData, internalSelectedURI, azureRows, loadingStatus]);
+  }, [resourcePickerData, internalSelectedURI, rows, loadingStatus]);
 
   // Map the selected item into an array of rows
   const selectedResourceRows = useMemo(() => {
-    const found = internalSelectedURI && findRow(azureRows, internalSelectedURI);
+    const found = internalSelectedURI && findRow(rows, internalSelectedURI);
 
     return found
       ? [
@@ -96,34 +72,28 @@ const ResourcePicker = ({
           },
         ]
       : [];
-  }, [internalSelectedURI, azureRows]);
+  }, [internalSelectedURI, rows]);
 
   // Request resources for a expanded resource group
   const requestNestedRows = useCallback(
-    async (resourceGroupOrSubscription: ResourceRow) => {
+    async (parentRow: ResourceRow) => {
       // clear error message (also when loading cached resources)
       setErrorMessage(undefined);
 
       // If we already have children, we don't need to re-fetch them.
-      if (resourceGroupOrSubscription.children?.length) {
+      if (parentRow.children?.length) {
         return;
       }
 
       try {
-        const rows =
-          resourceGroupOrSubscription.type === ResourceRowType.Subscription
-            ? await resourcePickerData.getResourceGroupsBySubscriptionId(resourceGroupOrSubscription.id)
-            : await resourcePickerData.getResourcesForResourceGroup(resourceGroupOrSubscription.id);
-
-        const newRows = addResources(azureRows, resourceGroupOrSubscription.uri, rows);
-
-        setAzureRows(newRows);
+        const nestedRows = await resourcePickerData.fetchAndAppendNestedRow(rows, parentRow);
+        setRows(nestedRows);
       } catch (error) {
         setErrorMessage(messageFromError(error));
         throw error;
       }
     },
-    [resourcePickerData, azureRows]
+    [resourcePickerData, rows]
   );
 
   const handleSelectionChanged = useCallback((row: ResourceRow, isSelected: boolean) => {
@@ -155,7 +125,7 @@ const ResourcePicker = ({
           <div className={styles.tableScroller}>
             <table className={styles.table}>
               <tbody>
-                {azureRows.map((row) => (
+                {rows.map((row) => (
                   <NestedRow
                     key={row.uri}
                     row={row}

public/app/plugins/datasource/grafana-azure-monitor-datasource/resourcePicker/resourcePickerData.ts

@@ -7,8 +7,8 @@ import {
   logsSupportedResourceTypesKusto,
   resourceTypeDisplayNames,
 } from '../azureMetadata';
-import { ResourceRowGroup, ResourceRowType } from '../components/ResourcePicker/types';
-import { parseResourceURI } from '../components/ResourcePicker/utils';
+import { ResourceRow, ResourceRowGroup, ResourceRowType } from '../components/ResourcePicker/types';
+import { addResources, parseResourceURI } from '../components/ResourcePicker/utils';
 import {
   AzureDataSourceJsonData,
   AzureGraphResponse,
@@ -31,6 +31,40 @@ export default class ResourcePickerData extends DataSourceWithBackend<AzureMonit
     this.resourcePath = `${routeNames.resourceGraph}`;
   }
 
+  async fetchInitialRows(currentSelection?: string): Promise<ResourceRowGroup> {
+    const subscriptions = await this.getSubscriptions();
+    if (!currentSelection) {
+      return subscriptions;
+    }
+
+    let resources = subscriptions;
+    const parsedURI = parseResourceURI(currentSelection);
+    if (parsedURI) {
+      const resourceGroupURI = `/subscriptions/${parsedURI.subscriptionID}/resourceGroups/${parsedURI.resourceGroup}`;
+
+      if (parsedURI.resourceGroup) {
+        const resourceGroups = await this.getResourceGroupsBySubscriptionId(parsedURI.subscriptionID);
+        resources = addResources(resources, `/subscriptions/${parsedURI.subscriptionID}`, resourceGroups);
+      }
+
+      if (parsedURI.resource) {
+        const resourcesForResourceGroup = await this.getResourcesForResourceGroup(resourceGroupURI);
+        resources = addResources(resources, resourceGroupURI, resourcesForResourceGroup);
+      }
+    }
+    return resources;
+  }
+
+  async fetchAndAppendNestedRow(rows: ResourceRowGroup, parentRow: ResourceRow): Promise<ResourceRowGroup> {
+    const nestedRows =
+      parentRow.type === ResourceRowType.Subscription
+        ? await this.getResourceGroupsBySubscriptionId(parentRow.id)
+        : await this.getResourcesForResourceGroup(parentRow.id);
+
+    return addResources(rows, parentRow.uri, nestedRows);
+  }
+
+  // private
   async getSubscriptions(): Promise<ResourceRowGroup> {
     const query = `
     resources