Fix jobs when used with cmdlets that require parametersetname

src/Common/Commands.Common/AzureLongRunningJob.cs

@@ -26,7 +26,59 @@
 
 namespace Microsoft.Azure.Commands.Common
 {
-    public class AzureLongRunningJob<T> : Job, ICommandRuntime where T : AzurePSCmdlet
+    /// <summary>
+    /// Abstract class for uniform processing of jobs
+    /// </summary>
+    public abstract class AzureLongRunningJob : Job
+    {
+        protected AzureLongRunningJob(string command, string name) : base(command, name)
+        {
+        }
+
+        /// <summary>
+        /// Run a cmdlet execution as a Job - this is intended to run in a background thread
+        /// </summary>
+        /// <param name="state">The Job record that will track the progress of this job </param>
+        public abstract void RunJob(object state);
+
+        /// <summary>
+        /// Return cmdlet runtime that will report results to this job
+        /// </summary>
+        /// <returns>An IcommandRuntime that reports results to this job</returns>
+        public abstract ICommandRuntime GetJobRuntime();
+
+        /// <summary>
+        /// Mark the job as started
+        /// </summary>
+        public abstract bool TryStart();
+
+        /// <summary>
+        /// Mark the job as Blocked
+        /// </summary>
+        public abstract bool TryBlock();
+
+
+        /// <summary>
+        /// Complete the job (will mark job as Completed or Failed, depending on the execution details)
+        /// </summary>
+        public abstract void Complete();
+
+        /// <summary>
+        /// Mark the Job as Failed
+        /// </summary>
+        public abstract void Fail();
+
+        /// <summary>
+        /// Stop the job
+        /// </summary>
+        public abstract void Cancel();
+    }
+
+    /// <summary>
+    /// Cmdlet-specific Implementation class for long running jobs
+    /// </summary>
+    /// <typeparam name="T">The type of the cmdlet being executed</typeparam>
+    public class AzureLongRunningJob<T> : AzureLongRunningJob, ICommandRuntime where T : AzurePSCmdlet
     {
         string _status = "Running";
         T _cmdlet;
@@ -196,14 +248,15 @@ public static U CopyCmdlet<U>(U cmdlet) where U : AzurePSCmdlet
                 field.SafeCopyValue(source: cmdlet, target: returnValue);
             }
 
+            cmdlet.SafeCopyParameterSet(returnValue);
             return returnValue as U;
         }
 
         /// <summary>
         /// Run a cmdlet execution as a Job - this is intended to run in a background thread
         /// </summary>
         /// <param name="state">The Job record that will track the progress of this job </param>
-        public virtual void RunJob(object state)
+        public override void RunJob(object state)
         {
             if (TryStart())
             {
@@ -239,7 +292,7 @@ public virtual void RunJob(object state)
         /// Return cmdlet runtime that will report results to this job
         /// </summary>
         /// <returns>An IcommandRuntime that reports results to this job</returns>
-        public virtual ICommandRuntime GetJobRuntime()
+        public override ICommandRuntime GetJobRuntime()
         {
             return this as ICommandRuntime;
         }
@@ -307,7 +360,7 @@ internal string GetShouldMethodFailureReaon(T cmdlet, ShouldMethodStreamItem met
         /// <summary>
         /// Mark the task as started
         /// </summary>
-        public bool TryStart()
+        public override bool TryStart()
         {
             bool result = false;
             lock (_lockObject)
@@ -326,7 +379,7 @@ public bool TryStart()
         /// <summary>
         /// Mark the task as blocked
         /// </summary>
-        public bool TryBlock()
+        public override bool TryBlock()
         {
             bool result = false;
             lock (_lockObject)
@@ -345,7 +398,7 @@ public bool TryBlock()
         /// <summary>
         /// Mark the job as Failed
         /// </summary>
-        public void Fail()
+        public override void Fail()
         {
             lock (_lockObject)
             {
@@ -357,7 +410,7 @@ public void Fail()
         /// <summary>
         /// Mark the job as successfully complete
         /// </summary>
-        public void Complete()
+        public override void Complete()
         {
             lock (_lockObject)
             {
@@ -378,7 +431,7 @@ public void Complete()
         /// <summary>
         /// Mark the job as cancelled
         /// </summary>
-        public void Cancel()
+        public override void Cancel()
         {
             lock (_lockObject)
             {
@@ -787,8 +840,7 @@ private bool InvokeShouldMethodAndWaitForResults(Func<Cmdlet, bool> shouldMethod
                 try
                 {
                     stateChangedEventHandler(null, new JobStateEventArgs(this.JobStateInfo));
-
-                    if (!gotResultEvent.IsSet && this.TryBlock())
+                    if (!gotResultEvent.IsSet)
                     {
                         WriteDebug(string.Format(Resources.TraceBlockLROThread, methodType));
                         ShouldMethodInvoker methodInvoker = new ShouldMethodInvoker
@@ -800,16 +852,18 @@ private bool InvokeShouldMethodAndWaitForResults(Func<Cmdlet, bool> shouldMethod
                         };
 
                         BlockedActions.Enqueue(new ShouldMethodStreamItem(methodInvoker));
-
-                        gotResultEvent.Wait();
-                        WriteDebug(string.Format(Resources.TraceUnblockLROThread, shouldMethod));
-                        TryStart();
-                        lock (resultsLock)
+                        if (this.TryBlock())
                         {
-                            if (closureSafeExceptionThrownOnCmdletThread == null) // stateChangedEventHandler didn't set the results?  = ok to clobber results?
+                            gotResultEvent.Wait();
+                            WriteDebug(string.Format(Resources.TraceUnblockLROThread, shouldMethod));
+                            TryStart();
+                            lock (resultsLock)
                             {
-                                closureSafeExceptionThrownOnCmdletThread = methodInvoker.ThrownException;
-                                methodResult = methodInvoker.MethodResult;
+                                if (closureSafeExceptionThrownOnCmdletThread == null) // stateChangedEventHandler didn't set the results?  = ok to clobber results?
+                                {
+                                    closureSafeExceptionThrownOnCmdletThread = methodInvoker.ThrownException;
+                                    methodResult = methodInvoker.MethodResult;
+                                }
                             }
                         }
                     }

src/Common/Commands.Common/Extensions/CmdletExtensions.cs

@@ -178,6 +178,30 @@ public static void SafeCopyValue<T>(this FieldInfo field, T source, T target)
             }
         }
 
+        /// <summary>
+        /// Safely copy the selected parameter set from one cmdlet to another
+        /// </summary>
+        /// <typeparam name="T">The cmdlet type</typeparam>
+        /// <param name="source">The cmdlet to copy the parameter set name from</param>
+        /// <param name="target">The cmdlet to copy to</param>
+        public static void SafeCopyParameterSet<T>(this T source, T target) where T : AzurePSCmdlet
+        {
+            if (source != null && target != null)
+            {
+                if (!string.IsNullOrWhiteSpace(source.ParameterSetName))
+                {
+                    try
+                    {
+                        target.SetParameterSet(source.ParameterSetName);
+                    }
+                    catch
+                    {
+
+                    }
+                }
+            }
+        }
+
         public static string AsAbsoluteLocation(this string realtivePath)
         {
             return Path.GetFullPath(Path.Combine(AppDomain.CurrentDomain.BaseDirectory, realtivePath));