Merge pull request Azure#484 from brahmnes/dev

Fixes to set and remove diag extension cmdlets

Author: ogail

src/ServiceManagement/Compute/Commands.ServiceManagement/IaaS/Extensions/Diagnostics/RemoveAzureVMDiagnosticsExtension.cs

@@ -30,11 +30,11 @@ public class RemoveAzureVMDiagnosticsExtensionCommand : VirtualMachineDiagnostic
 
         internal void ExecuteCommand()
         {
-            // We do not support this cmdlet for now...
-            throw new NotSupportedException();
+            ExtensionName = DiagnosticsExtensionType;
+            Publisher = DiagnosticsExtensionNamespace;
 
-        //    RemovePredicateExtensions();   
-        //    WriteObject(VM);
+            RemovePredicateExtensions();
+            WriteObject(VM);
         }
 
         protected override void ProcessRecord()

src/ServiceManagement/Compute/Commands.ServiceManagement/IaaS/Extensions/Diagnostics/SetAzureVMDiagnosticsExtension.cs

@@ -14,6 +14,7 @@
 
 using System;
 using System.IO;
+using System.Linq;
 using System.Management.Automation;
 using System.Text;
 using System.Xml;
@@ -121,6 +122,17 @@ protected override void ValidateParameters()
             ValidateConfiguration();
             ExtensionName = DiagnosticsExtensionType;
             Publisher = DiagnosticsExtensionNamespace;
+
+            // If the user didn't specify an extension reference name and the input VM already has a diagnostics extension,
+            // reuse its reference name
+            if (string.IsNullOrEmpty(ReferenceName))
+            {
+                ResourceExtensionReference diagnosticsExtension = ResourceExtensionReferences.FirstOrDefault(ExtensionPredicate);
+                if (diagnosticsExtension != null)
+                {
+                    ReferenceName = diagnosticsExtension.ReferenceName;
+                }
+            }
         }
 
         private void ValidateStorageAccount()