From 955405dd394d4ff47a742fd68682caaf213205b6 Mon Sep 17 00:00:00 2001
From: MicroFish91 <mwfisher91@gmail.com>
Date: Mon, 18 Dec 2023 17:41:46 -0600
Subject: [PATCH] Impl

---
 utils/index.d.ts                               |  6 +++++-
 utils/src/pickTreeItem/GenericQuickPickStep.ts | 12 +++++++++++-
 .../experiences/subscriptionExperience.ts      | 18 +++++++++++-------
 .../QuickPickAzureSubscriptionStep.ts          | 16 ++++++++++++++--
 4 files changed, 41 insertions(+), 11 deletions(-)

diff --git a/utils/index.d.ts b/utils/index.d.ts
index c210a0e9d6..20597bdb2d 100644
--- a/utils/index.d.ts
+++ b/utils/index.d.ts
@@ -1733,7 +1733,7 @@ export interface PickExperienceContext extends IActionContext {
  * @param context The action context
  * @param tdp Azure resource tree data provider to perform the pick experience on
  */
-export declare function subscriptionExperience(context: IActionContext, tdp: TreeDataProvider<ResourceGroupsItem>): Promise<AzureSubscription>;
+export declare function subscriptionExperience(context: IActionContext, tdp: TreeDataProvider<ResourceGroupsItem>, options?: { selectSubscriptionId?: string, showLoadingPrompt?: boolean }): Promise<AzureSubscription>;
 
 /**
  * Prompts the user to pick an Azure resource using a wizard comprised of quick pick steps.
@@ -1876,6 +1876,10 @@ export declare interface SkipIfOneQuickPickOptions extends GenericQuickPickOptio
     skipIfOne?: true;
 }
 
+export declare interface AzureSubscriptionQuickPickOptions extends GenericQuickPickOptions {
+    selectSubscriptionId?: string;
+}
+
 export declare abstract class GenericQuickPickStep<TContext extends QuickPickWizardContext, TOptions extends GenericQuickPickOptions> extends AzureWizardPromptStep<TContext> {
     constructor(
         treeDataProvider: TreeDataProvider<unknown>,
diff --git a/utils/src/pickTreeItem/GenericQuickPickStep.ts b/utils/src/pickTreeItem/GenericQuickPickStep.ts
index 48736a033f..abef1cb6f6 100644
--- a/utils/src/pickTreeItem/GenericQuickPickStep.ts
+++ b/utils/src/pickTreeItem/GenericQuickPickStep.ts
@@ -6,8 +6,8 @@
 import * as vscode from 'vscode';
 import * as types from '../../index';
 import { AzureWizardPromptStep } from '../wizard/AzureWizardPromptStep';
-import { getLastNode } from './getLastNode';
 import { PickFilter } from './PickFilter';
+import { getLastNode } from './getLastNode';
 
 export abstract class GenericQuickPickStep<TContext extends types.QuickPickWizardContext, TOptions extends types.GenericQuickPickOptions> extends AzureWizardPromptStep<TContext> {
     public readonly supportsDuplicateSteps = true;
@@ -50,6 +50,14 @@ export abstract class GenericQuickPickStep<TContext extends types.QuickPickWizar
             }
         }
 
+        const pickWithoutPrompt = this.getPickWithoutPrompt?.(picks);
+        if (pickWithoutPrompt) {
+            const ti = await this.treeDataProvider.getTreeItem(pickWithoutPrompt.data);
+            if (!ti.command) {
+                return pickWithoutPrompt.data;
+            }
+        }
+
         const selected = await wizardContext.ui.showQuickPick(picks, {
             ...(this.promptOptions ?? {})
         });
@@ -94,4 +102,6 @@ export abstract class GenericQuickPickStep<TContext extends types.QuickPickWizar
             data: element,
         };
     }
+
+    protected getPickWithoutPrompt?(picks: types.IAzureQuickPickItem<unknown>[]): types.IAzureQuickPickItem<unknown> | undefined;
 }
diff --git a/utils/src/pickTreeItem/experiences/subscriptionExperience.ts b/utils/src/pickTreeItem/experiences/subscriptionExperience.ts
index afac4719b2..1b06f46354 100644
--- a/utils/src/pickTreeItem/experiences/subscriptionExperience.ts
+++ b/utils/src/pickTreeItem/experiences/subscriptionExperience.ts
@@ -3,24 +3,28 @@
 *  Licensed under the MIT License. See License.txt in the project root for license information.
 *--------------------------------------------------------------------------------------------*/
 
-import * as types from '../../../index';
+import type { AzureSubscription } from '@microsoft/vscode-azureresources-api';
 import * as vscode from 'vscode';
+import * as types from '../../../index';
+import { PickSubscriptionWizardContext } from '../../../index';
+import { NoResourceFoundError } from '../../errors';
 import { AzureWizard } from '../../wizard/AzureWizard';
 import { QuickPickAzureSubscriptionStep } from '../quickPickAzureResource/QuickPickAzureSubscriptionStep';
-import { PickSubscriptionWizardContext } from '../../../index';
 import { ResourceGroupsItem } from '../quickPickAzureResource/tempTypes';
-import { NoResourceFoundError } from '../../errors';
-import type { AzureSubscription } from '@microsoft/vscode-azureresources-api';
 
-export async function subscriptionExperience(context: types.IActionContext, tdp: vscode.TreeDataProvider<ResourceGroupsItem>): Promise<AzureSubscription> {
+export async function subscriptionExperience(
+    context: types.IActionContext,
+    tdp: vscode.TreeDataProvider<ResourceGroupsItem>,
+    options?: { selectSubscriptionId?: string, showLoadingPrompt?: boolean }
+): Promise<AzureSubscription> {
 
     const wizardContext = { ...context } as PickSubscriptionWizardContext;
     wizardContext.pickedNodes = [];
 
     const wizard = new AzureWizard(wizardContext, {
         hideStepCount: true,
-        promptSteps: [new QuickPickAzureSubscriptionStep(tdp)],
-        showLoadingPrompt: true,
+        promptSteps: [new QuickPickAzureSubscriptionStep(tdp, { selectSubscriptionId: options?.selectSubscriptionId })],
+        showLoadingPrompt: options?.showLoadingPrompt ?? true,
     });
 
     await wizard.prompt();
diff --git a/utils/src/pickTreeItem/quickPickAzureResource/QuickPickAzureSubscriptionStep.ts b/utils/src/pickTreeItem/quickPickAzureResource/QuickPickAzureSubscriptionStep.ts
index 20f941a612..33db8956a5 100644
--- a/utils/src/pickTreeItem/quickPickAzureResource/QuickPickAzureSubscriptionStep.ts
+++ b/utils/src/pickTreeItem/quickPickAzureResource/QuickPickAzureSubscriptionStep.ts
@@ -4,13 +4,15 @@
  *--------------------------------------------------------------------------------------------*/
 
 import * as vscode from 'vscode';
-import { AzureResourceQuickPickWizardContext, GenericQuickPickOptions, SkipIfOneQuickPickOptions } from '../../../index';
+import { AzureResourceQuickPickWizardContext, AzureSubscriptionQuickPickOptions, IAzureQuickPickItem, SkipIfOneQuickPickOptions } from '../../../index';
 import { GenericQuickPickStepWithCommands } from '../GenericQuickPickStepWithCommands';
 import { PickFilter } from '../PickFilter';
 import { ResourceGroupsItem, SubscriptionItem } from './tempTypes';
 
 export class QuickPickAzureSubscriptionStep extends GenericQuickPickStepWithCommands<AzureResourceQuickPickWizardContext, SkipIfOneQuickPickOptions> {
-    public constructor(tdp: vscode.TreeDataProvider<ResourceGroupsItem>, options?: GenericQuickPickOptions) {
+    protected subscriptionId?: string;
+
+    public constructor(tdp: vscode.TreeDataProvider<ResourceGroupsItem>, options?: AzureSubscriptionQuickPickOptions) {
         super(tdp, {
             ...options,
             skipIfOne: true, // Subscription is always skip-if-one
@@ -18,6 +20,8 @@ export class QuickPickAzureSubscriptionStep extends GenericQuickPickStepWithComm
             placeHolder: vscode.l10n.t('Select subscription'),
             noPicksMessage: vscode.l10n.t('No subscriptions found'),
         });
+
+        this.subscriptionId = options?.selectSubscriptionId;
     }
 
     protected readonly pickFilter = new AzureSubscriptionPickFilter();
@@ -30,6 +34,14 @@ export class QuickPickAzureSubscriptionStep extends GenericQuickPickStepWithComm
 
         return pickedSubscription;
     }
+
+    protected override getPickWithoutPrompt(picks: IAzureQuickPickItem<SubscriptionItem>[]): IAzureQuickPickItem<SubscriptionItem> | undefined {
+        if (!this.subscriptionId) {
+            return undefined;
+        }
+
+        return picks.find(pick => pick.data.subscription.subscriptionId === this.subscriptionId);
+    }
 }
 
 class AzureSubscriptionPickFilter implements PickFilter {