From 284afbcd4b9c4b16ddc70a3cf8b6d831369d6c70 Mon Sep 17 00:00:00 2001
From: Alyssa Bull <58453566+alyssabull@users.noreply.github.com>
Date: Wed, 10 Jan 2024 13:58:03 -0700
Subject: [PATCH] Azure Monitor: Add select all subscription option for ARG
 queries (#79582)

---
 .../ArgQueryEditor/ArgQueryEditor.test.tsx    | 51 +++++++++++++++++++
 .../ArgQueryEditor/ArgQueryEditor.tsx         |  5 +-
 .../ArgQueryEditor/SubscriptionField.tsx      |  7 +++
 3 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/public/app/plugins/datasource/azuremonitor/components/ArgQueryEditor/ArgQueryEditor.test.tsx b/public/app/plugins/datasource/azuremonitor/components/ArgQueryEditor/ArgQueryEditor.test.tsx
index 3f5f31a1966..92564c4c7dc 100644
--- a/public/app/plugins/datasource/azuremonitor/components/ArgQueryEditor/ArgQueryEditor.test.tsx
+++ b/public/app/plugins/datasource/azuremonitor/components/ArgQueryEditor/ArgQueryEditor.test.tsx
@@ -182,4 +182,55 @@ describe('ArgQueryEditor', () => {
     );
     expect(await waitFor(() => screen.getByText('At least one subscription must be chosen.'))).toBeInTheDocument();
   });
+
+  it('should select all subscriptions if select all is chosen from the dropdown', async () => {
+    const onChange = jest.fn();
+    const datasource = createMockDatasource({
+      getSubscriptions: jest.fn().mockResolvedValue([
+        { text: 'foo', value: 'test-subscription-value1' },
+        { text: 'bar', value: 'test-subscription-value2' },
+        { text: 'Select all subscriptions', value: 'Select all' },
+      ]),
+    });
+    const query = createMockQuery({
+      subscription: undefined,
+      subscriptions: ['test-subscription-value1', 'test-subscription-value2', 'Select all'],
+    });
+    const { rerender } = render(
+      <ArgQueryEditor
+        {...defaultProps}
+        query={query}
+        datasource={datasource}
+        onChange={onChange}
+        variableOptionGroup={{ label: 'Template Variables', options: [] }}
+      />
+    );
+
+    expect(datasource.getSubscriptions).toHaveBeenCalled();
+    expect(await waitFor(() => onChange)).toHaveBeenCalledWith(
+      expect.objectContaining({ subscriptions: ['test-subscription-value1', 'test-subscription-value2', 'Select all'] })
+    );
+    expect(await waitFor(() => screen.findByText('foo'))).toBeInTheDocument();
+    expect(await waitFor(() => screen.findByText('bar'))).toBeInTheDocument();
+    expect(await waitFor(() => screen.findByText('Select all subscriptions'))).toBeInTheDocument();
+
+    const selectAll = screen.getByText('Select all subscriptions');
+    await userEvent.click(selectAll);
+
+    expect(onChange).toHaveBeenCalledWith(
+      expect.objectContaining({ subscriptions: ['test-subscription-value1', 'test-subscription-value2', 'Select all'] })
+    );
+
+    rerender(
+      <ArgQueryEditor
+        {...defaultProps}
+        datasource={datasource}
+        onChange={onChange}
+        query={{ ...query, subscriptions: ['test-subscription-value1', 'test-subscription-value2', 'Select all'] }}
+        variableOptionGroup={{ label: 'Template Variables', options: [] }}
+      />
+    );
+    expect(await waitFor(() => screen.getByText('foo'))).toBeInTheDocument();
+    expect(await waitFor(() => screen.getByText('bar'))).toBeInTheDocument();
+  });
 });
diff --git a/public/app/plugins/datasource/azuremonitor/components/ArgQueryEditor/ArgQueryEditor.tsx b/public/app/plugins/datasource/azuremonitor/components/ArgQueryEditor/ArgQueryEditor.tsx
index 78b685d491b..cc0d1ce82d9 100644
--- a/public/app/plugins/datasource/azuremonitor/components/ArgQueryEditor/ArgQueryEditor.tsx
+++ b/public/app/plugins/datasource/azuremonitor/components/ArgQueryEditor/ArgQueryEditor.tsx
@@ -57,8 +57,11 @@ const ArgQueryEditor = ({
     datasource
       .getSubscriptions()
       .then((results) => {
+        const selectAllSubscriptionOption = [
+          { label: 'Select all subscriptions', value: 'Select all subscriptions', description: 'Select all' },
+        ];
         const fetchedSubscriptions = results.map((v) => ({ label: v.text, value: v.value, description: v.value }));
-        setSubscriptions(fetchedSubscriptions);
+        setSubscriptions(selectAllSubscriptionOption.concat(fetchedSubscriptions));
         setError(ERROR_SOURCE, undefined);
 
         onChange({
diff --git a/public/app/plugins/datasource/azuremonitor/components/ArgQueryEditor/SubscriptionField.tsx b/public/app/plugins/datasource/azuremonitor/components/ArgQueryEditor/SubscriptionField.tsx
index 27b3a950034..7b03af248bc 100644
--- a/public/app/plugins/datasource/azuremonitor/components/ArgQueryEditor/SubscriptionField.tsx
+++ b/public/app/plugins/datasource/azuremonitor/components/ArgQueryEditor/SubscriptionField.tsx
@@ -28,6 +28,7 @@ const SubscriptionField = ({ query, subscriptions, variableOptionGroup, onQueryC
   }, [query.subscriptions, subscriptions, variableOptionGroup.options]);
 
   const onChange = (change: Array<SelectableValue<string>>) => {
+    const containsSelectAll = change.filter((c) => c.value === 'Select all subscriptions');
     if (!change || change.length === 0) {
       setValues([]);
       onQueryChange({
@@ -35,6 +36,12 @@ const SubscriptionField = ({ query, subscriptions, variableOptionGroup, onQueryC
         subscriptions: [],
       });
       setError(true);
+    } else if (containsSelectAll.length > 0) {
+      const allSubs = subscriptions.map((c) => c.value ?? '').filter((c) => c !== 'Select all subscriptions');
+      onQueryChange({
+        ...query,
+        subscriptions: allSubs,
+      });
     } else {
       const newSubs = change.map((c) => c.value ?? '');
       onQueryChange({