Merge pull request #6004 from maddieclayton/AsJobBug

Fix Issue 5896

Author: cormacpayne

src/Common/Commands.Common/AzurePSCmdlet.cs

@@ -137,13 +137,50 @@ protected string PSVersion
 
         protected abstract string DataCollectionWarning { get; }
 
+        private SessionState _sessionState;
+
+        public new SessionState SessionState
+        {
+            get
+            {
+                return _sessionState;
+            }
+            set
+            {
+                _sessionState = value;
+            }
+        }
+
+        private RuntimeDefinedParameterDictionary _asJobDynamicParameters;
+
+        public RuntimeDefinedParameterDictionary AsJobDynamicParameters
+        {
+            get
+            {
+                if (_asJobDynamicParameters == null)
+                {
+                    _asJobDynamicParameters = new RuntimeDefinedParameterDictionary();
+                }
+                return _asJobDynamicParameters;
+            }
+            set
+            {
+                _asJobDynamicParameters = value;
+            }
+        }
+
         /// <summary>
         /// Initializes AzurePSCmdlet properties.
         /// </summary>
         public AzurePSCmdlet()
         {
             DebugMessages = new ConcurrentQueue<string>();
+        }
 
+        // Register Dynamic Parameters for use in long running jobs
+        public void RegisterDynamicParameters(RuntimeDefinedParameterDictionary parameters)
+        {
+            this.AsJobDynamicParameters = parameters;
         }
 
 
@@ -252,6 +289,7 @@ protected virtual void TearDownHttpClientPipeline()
         /// </summary>
         protected override void BeginProcessing()
         {
+            SessionState = base.SessionState;
             var profile = _dataCollectionProfile;
             //TODO: Inject from CI server
             lock (lockObject)

src/ResourceManager/Compute/Commands.Compute.Test/ScenarioTests/VirtualMachineRunCommandTests.ps1

@@ -107,7 +107,9 @@ function Test-VirtualMachineSetRunCommand
 
         $path = 'ScenarioTests\test.ps1';
 
-        $result = Invoke-AzureRmVMRunCommand -ResourceGroupName $rgname -Name $vmname -CommandId $commandId -ScriptPath $path -Parameter $param;
+        $job = Invoke-AzureRmVMRunCommand -ResourceGroupName $rgname -Name $vmname -CommandId $commandId -ScriptPath $path -Parameter $param -AsJob;
+		$result = $job | Wait-Job
+
         Assert-AreEqual "Succeeded" $result.Status
 
         # Remove All VMs

src/ResourceManager/Profile/Commands.Profile.Test/LongRunningCmdletTests.cs

@@ -183,6 +183,7 @@ public void JobCopiesCmdletParameterSet()
             Mock<ICommandRuntime> mock = new Mock<ICommandRuntime>();
             var cmdlet = new AzureParameterSetCmdlet();
             cmdlet.SetParameterSet("ParameterSetIsSet");
+            cmdlet.AsJobDynamicParameters = new RuntimeDefinedParameterDictionary();
             cmdlet.CommandRuntime = mock.Object;
             var job = cmdlet.ExecuteAsJob("Test parameter set job") as AzureLongRunningJob<AzureParameterSetCmdlet>;
             WaitForCompletion(job, j =>
@@ -389,6 +390,10 @@ public override void ExecuteCmdlet()
                 {
                     throw new InvalidOperationException("Parameter set must be set");
                 }
+                if (this.AsJobDynamicParameters == null)
+                {
+                    throw new InvalidOperationException("Dynamic parameters must be set");
+                }
             }
 
             protected override void InitializeQosEvent()

src/ResourceManager/Resources/Commands.ResourceManager/Cmdlets/Implementation/CmdletBase/ResourceWithParameterCmdletBase.cs

@@ -135,6 +135,8 @@ public object GetDynamicParameters()
                 }
             }
 
+            RegisterDynamicParameters(dynamicParameters);
+
             return dynamicParameters;
         }
 
@@ -172,7 +174,7 @@ protected Hashtable GetTemplateParameterObject(Hashtable templateParameterObject
             }
 
             // Load dynamic parameters
-            IEnumerable<RuntimeDefinedParameter> parameters = PowerShellUtilities.GetUsedDynamicParameters(dynamicParameters, MyInvocation);
+            IEnumerable<RuntimeDefinedParameter> parameters = PowerShellUtilities.GetUsedDynamicParameters(this.AsJobDynamicParameters, MyInvocation);
             if (parameters.Any())
             {
                 parameters.ForEach(dp => prameterObject[((ParameterAttribute)dp.Attributes[0]).HelpMessage] = new Hashtable { { "value", dp.Value } });

src/ResourceManager/Resources/Commands.ResourceManager/Cmdlets/Implementation/Policy/NewAzurePolicyAssignment.cs

@@ -248,6 +248,8 @@ object IDynamicParameters.GetDynamicParameters()
                 }
             }
 
+            RegisterDynamicParameters(dynamicParameters);
+
             return this.dynamicParameters;
         }
 
@@ -269,7 +271,7 @@ private JObject GetParameters()
             }
 
             // Load dynamic parameters
-            var parameters = PowerShellUtilities.GetUsedDynamicParameters(dynamicParameters, MyInvocation);
+            var parameters = PowerShellUtilities.GetUsedDynamicParameters(AsJobDynamicParameters, MyInvocation);
             if (parameters.Count() > 0)
             {
                 return MyInvocation.BoundParameters.ToJObjectWithValue(parameters.Select(p => p.Name));