Removed handler and moved pre-register logic to cmdlet

Author: Sean Oh

src/ResourceManager/Reservations/Commands.Reservations/Cmdlets/PatchReservation.cs

@@ -6,11 +6,17 @@
 using System.Management.Automation;
 using Newtonsoft.Json;
 using Microsoft.Azure.Management.Reservations;
+using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
+using Microsoft.Azure.Management.Internal.Resources;
+using Microsoft.Azure.Commands.Common.Authentication;
+using Microsoft.WindowsAzure.Commands.Utilities.Common;
+using Microsoft.Azure.Management.Internal.Resources.Models;
+using Microsoft.Rest.Azure;
 
 namespace Microsoft.Azure.Commands.Reservations.Cmdlets
 {
     [Cmdlet(VerbsData.Update, "AzureRmReservation", DefaultParameterSetName = Constants.ParameterSetNames.CommandParameterSet, SupportsShouldProcess = true), OutputType(typeof(PSReservation))]
-    public class PatchReservation : AzureReservationsAutoRegisterCmdletBase
+    public class PatchReservation : AzureReservationsCmdletBase
     {
         [Parameter(ParameterSetName = Constants.ParameterSetNames.CommandParameterSet,
             Mandatory = true)]
@@ -37,8 +43,6 @@ public class PatchReservation : AzureReservationsAutoRegisterCmdletBase
         [ValidateNotNull]
         public PSReservation Reservation { get; set; }
 
-        protected override bool ShouldRegister => AppliedScope != null;
-
         public override void ExecuteCmdlet()
         {
             if (ParameterSetName.Equals(Constants.ParameterSetNames.ObjectParameterSet))
@@ -55,6 +59,10 @@ public override void ExecuteCmdlet()
                 Patch Patch;
                 if (AppliedScope != null)
                 {
+                    //Pre-register for Microsoft.Compute
+                    string subscriptionId = ValidateAndGetAppliedSubscription();
+                    PreRegister(subscriptionId);
+
                     Patch = new Patch(AppliedScopeType, new List<string>() { AppliedScope });
                 }
                 else
@@ -65,5 +73,68 @@ public override void ExecuteCmdlet()
                 WriteObject(response);
             }
         }
+
+        private void PreRegister(string subscriptionId)
+        {
+            try
+            {
+                IAzureContext context;
+                if (TryGetDefaultContext(out context)
+                    && context.Account != null
+                    && context.Subscription != null)
+                {
+                    var client = new ResourceManagementClient(
+                                context.Environment.GetEndpointAsUri(AzureEnvironment.Endpoint.ResourceManager),
+                                AzureSession.Instance.AuthenticationFactory.GetServiceClientCredentials(context, AzureEnvironment.Endpoint.ResourceManager));
+                    client.SubscriptionId = subscriptionId;
+
+                    string ComputeProviderNamespace = "Microsoft.Compute";
+                    var maxRetryCount = 10;
+
+                    var provider = client.Providers.Get(ComputeProviderNamespace);
+                    if (provider.RegistrationState != RegistrationState.Registered)
+                    {
+                        short retryCount = 0;
+                        do
+                        {
+                            if (retryCount++ > maxRetryCount)
+                            {
+                                throw new TimeoutException();
+                            }
+                            provider = client.Providers.Register(ComputeProviderNamespace);
+                            TestMockSupport.Delay(2000);
+                        } while (provider.RegistrationState != RegistrationState.Registered);
+                    }
+                }
+            }
+            catch (Exception e)
+            {
+                if (e.Message?.IndexOf("does not have authorization") >= 0 && e.Message?.IndexOf("register/action",
+                        StringComparison.InvariantCultureIgnoreCase) >= 0)
+                {
+                    throw new CloudException(e.Message);
+                }
+            }
+        }
+
+        private string ValidateAndGetAppliedSubscription()
+        {
+            string subscriptionId = AppliedScope;
+            string prefix = "/subscriptions/";
+            if (subscriptionId.Contains(prefix) && subscriptionId.Length > prefix.Length)
+            {
+                subscriptionId = subscriptionId.Substring(prefix.Length);
+            }
+
+            Guid result;
+            if (Guid.TryParse(subscriptionId, out result))
+            {
+                return result.ToString();
+            }
+            else
+            {
+                throw new PSArgumentException("Invalid applied scope provided");
+            }
+        }
     }
 }

src/ResourceManager/Reservations/Commands.Reservations/Commands.Reservations.csproj

@@ -62,9 +62,7 @@
     <Compile Include="Cmdlets\MergeReservation.cs" />
     <Compile Include="Cmdlets\PatchReservation.cs" />
     <Compile Include="Cmdlets\SplitReservation.cs" />
-    <Compile Include="Common\AzureReservationsAutoRegisterCmdletBase.cs" />
     <Compile Include="Common\AzureReservationsCmdletBase.cs" />
-    <Compile Include="Common\AzureReservationsAutoRegisterDelegatingHandler.cs" />
     <Compile Include="Common\Constants.cs" />
     <Compile Include="Custom\PSCatalog.cs" />
     <Compile Include="Custom\PSAppliedReservationOrderId.cs" />