From bdf59d5fcad5014b5d703938eb3c25338f1efd8e Mon Sep 17 00:00:00 2001
From: Weidong Xu <weidxu@microsoft.com>
Date: Thu, 4 Mar 2021 14:30:38 +0800
Subject: [PATCH] support withBootDiagnosticsOnManagedStorageAccount

---
 .../management/compute/VirtualMachine.java    | 16 +++++++++++++++
 .../implementation/VirtualMachineImpl.java    | 20 +++++++++++++++++--
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/azure-mgmt-compute/src/main/java/com/microsoft/azure/management/compute/VirtualMachine.java b/azure-mgmt-compute/src/main/java/com/microsoft/azure/management/compute/VirtualMachine.java
index 2908bb90790..fa011a7a880 100644
--- a/azure-mgmt-compute/src/main/java/com/microsoft/azure/management/compute/VirtualMachine.java
+++ b/azure-mgmt-compute/src/main/java/com/microsoft/azure/management/compute/VirtualMachine.java
@@ -1650,6 +1650,14 @@ interface WithPlan {
          * The stage of the virtual machine definition allowing to enable boot diagnostics.
          */
         interface WithBootDiagnostics {
+            /**
+             * Specifies that boot diagnostics needs to be enabled in the virtual machine.
+             * Managed storage account is used.
+             *
+             * @return the next stage of the definition
+             */
+            WithCreate withBootDiagnosticsOnManagedStorageAccount();
+
             /**
              * Specifies that boot diagnostics needs to be enabled in the virtual machine.
              *
@@ -2221,6 +2229,14 @@ interface WithExtension {
          * The stage of the virtual machine definition allowing to enable boot diagnostics.
          */
         interface WithBootDiagnostics {
+            /**
+             * Specifies that boot diagnostics needs to be enabled in the virtual machine.
+             * Managed storage account is used.
+             *
+             * @return the next stage of the update
+             */
+            Update withBootDiagnosticsOnManagedStorageAccount();
+
             /**
              * Specifies that boot diagnostics needs to be enabled in the virtual machine.
              *
diff --git a/azure-mgmt-compute/src/main/java/com/microsoft/azure/management/compute/implementation/VirtualMachineImpl.java b/azure-mgmt-compute/src/main/java/com/microsoft/azure/management/compute/implementation/VirtualMachineImpl.java
index b67b6fc76eb..aa410eeb6c9 100644
--- a/azure-mgmt-compute/src/main/java/com/microsoft/azure/management/compute/implementation/VirtualMachineImpl.java
+++ b/azure-mgmt-compute/src/main/java/com/microsoft/azure/management/compute/implementation/VirtualMachineImpl.java
@@ -1391,9 +1391,15 @@ public VirtualMachineImpl withUnmanagedDisks() {
         return this;
     }
 
+    @Override
+    public VirtualMachineImpl withBootDiagnosticsOnManagedStorageAccount() {
+        this.bootDiagnosticsHandler.withBootDiagnostics(true);
+        return this;
+    }
+
     @Override
     public VirtualMachineImpl withBootDiagnostics() {
-        this.bootDiagnosticsHandler.withBootDiagnostics();
+        this.bootDiagnosticsHandler.withBootDiagnostics(false);
         return this;
     }
 
@@ -2733,6 +2739,7 @@ private StorageAccountTypes getDefaultStorageAccountType() {
     private class BootDiagnosticsHandler {
         private final VirtualMachineImpl vmImpl;
         private String creatableDiagnosticsStorageAccountKey;
+        private boolean useManagedStorageAccount = false;
 
         BootDiagnosticsHandler(VirtualMachineImpl vmImpl) {
             this.vmImpl = vmImpl;
@@ -2756,10 +2763,11 @@ public String bootDiagnosticsStorageUri() {
             return null;
         }
 
-        BootDiagnosticsHandler withBootDiagnostics() {
+        BootDiagnosticsHandler withBootDiagnostics(boolean useManagedStorageAccount) {
             // Diagnostics storage uri will be set later by this.handleDiagnosticsSettings(..)
             //
             this.enableDisable(true);
+            this.useManagedStorageAccount = useManagedStorageAccount;
             return this;
         }
 
@@ -2773,6 +2781,7 @@ BootDiagnosticsHandler withBootDiagnostics(Creatable<StorageAccount> creatable)
 
         BootDiagnosticsHandler withBootDiagnostics(String storageAccountBlobEndpointUri) {
             this.enableDisable(true);
+            this.useManagedStorageAccount = false;
             this.vmInner()
                     .diagnosticsProfile()
                     .bootDiagnostics()
@@ -2786,10 +2795,14 @@ BootDiagnosticsHandler withBootDiagnostics(StorageAccount storageAccount) {
 
         BootDiagnosticsHandler withoutBootDiagnostics() {
             this.enableDisable(false);
+            this.useManagedStorageAccount = false;
             return this;
         }
 
         void prepare() {
+            if (useManagedStorageAccount) {
+                return;
+            }
             DiagnosticsProfile diagnosticsProfile = this.vmInner().diagnosticsProfile();
             if (diagnosticsProfile == null
                     || diagnosticsProfile.bootDiagnostics() == null
@@ -2822,6 +2835,9 @@ void prepare() {
         }
 
         void handleDiagnosticsSettings() {
+            if (useManagedStorageAccount) {
+                return;
+            }
             DiagnosticsProfile diagnosticsProfile = this.vmInner().diagnosticsProfile();
             if (diagnosticsProfile == null
                     || diagnosticsProfile.bootDiagnostics() == null