support withBootDiagnosticsOnManagedStorageAccount (#1372)

Author: weidongxu-microsoft

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.
              *

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