Add support for multiple compute nodes

Author: jasper-schneider

src/ResourceManager/AzureBatch/Commands.Batch.Test/Commands.Batch.Test.csproj

@@ -371,6 +371,9 @@
     <None Include="SessionRecords\Microsoft.Azure.Commands.Batch.Test.ScenarioTests.ComputeNodeTests\TestRemoveComputeNodePipeline.json">
       <CopyToOutputDirectory>Always</CopyToOutputDirectory>
     </None>
+    <None Include="SessionRecords\Microsoft.Azure.Commands.Batch.Test.ScenarioTests.ComputeNodeTests\TestRemoveMultipleComputeNodes.json">
+      <CopyToOutputDirectory>Always</CopyToOutputDirectory>
+    </None>
     <None Include="SessionRecords\Microsoft.Azure.Commands.Batch.Test.ScenarioTests.ComputeNodeUserTests\TestCreateComputeNodeUser.json">
       <CopyToOutputDirectory>Always</CopyToOutputDirectory>
     </None>

src/ResourceManager/AzureBatch/Commands.Batch.Test/ComputeNodes/RemoveBatchComputeNodeCommandTests.cs

@@ -54,15 +54,15 @@ public void RemoveBatchComputeNodeParametersTest()
 
             BatchAccountContext context = BatchTestHelpers.CreateBatchContextWithKeys();
             cmdlet.BatchContext = context;
-            cmdlet.Id = null;
+            cmdlet.Ids = null;
 
             Assert.Throws<ArgumentNullException>(() => cmdlet.ExecuteCmdlet());
 
             cmdlet.PoolId = "testPool";
 
             Assert.Throws<ArgumentNullException>(() => cmdlet.ExecuteCmdlet());
 
-            cmdlet.Id = "computeNode1";
+            cmdlet.Ids = new string[] { "computeNode1" };
 
             // Don't go to the service on a Remove ComputeNode call
             RequestInterceptor interceptor = BatchTestHelpers.CreateFakeServiceResponseInterceptor<ComputeNodeRemoveParameters, ComputeNodeRemoveResponse>();
@@ -84,12 +84,13 @@ public void RemoveComputeNodeRequestTest()
             cmdlet.BatchContext = context;
 
             cmdlet.PoolId = "testPool";
-            cmdlet.Id = "computeNode1";
+            cmdlet.Ids = new string[] { "computeNode1", "computeNode2" };
             cmdlet.DeallocationOption = ComputeNodeDeallocationOption.Terminate;
             cmdlet.ResizeTimeout = TimeSpan.FromMinutes(8);
 
             ComputeNodeDeallocationOption? requestDeallocationOption = null;
             TimeSpan? requestResizeTimeout = null;
+            IList<string> requestComputeNodeIds = null;
 
             // Don't go to the service on a Remove ComputeNode call
             RequestInterceptor interceptor = new RequestInterceptor((baseRequest) =>
@@ -102,6 +103,7 @@ public void RemoveComputeNodeRequestTest()
                     // Grab the parameters from the outgoing request.
                     requestDeallocationOption = request.TypedParameters.ComputeNodeDeallocationOption;
                     requestResizeTimeout = request.TypedParameters.ResizeTimeout;
+                    requestComputeNodeIds = request.TypedParameters.ComputeNodeIds;
 
                     ComputeNodeRemoveResponse response = new ComputeNodeRemoveResponse();
                     Task<ComputeNodeRemoveResponse> task = Task.FromResult(response);
@@ -115,6 +117,11 @@ public void RemoveComputeNodeRequestTest()
             // Verify that the parameters were properly set on the outgoing request
             Assert.Equal(cmdlet.DeallocationOption, requestDeallocationOption);
             Assert.Equal(cmdlet.ResizeTimeout, requestResizeTimeout);
+            Assert.Equal(cmdlet.Ids.Length, requestComputeNodeIds.Count);
+            foreach (string id in cmdlet.Ids)
+            {
+                Assert.True(requestComputeNodeIds.Contains(id));
+            }
         }
     }
 }

src/ResourceManager/AzureBatch/Commands.Batch.Test/ScenarioTests/ComputeNodeTests.cs

@@ -137,6 +137,34 @@ public void TestRemoveComputeNodePipeline()
             TestRemoveComputeNode(true, TestUtilities.GetCurrentMethodName());
         }
 
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void TestRemoveMultipleComputeNodes()
+        {
+            BatchController controller = BatchController.NewInstance;
+            BatchAccountContext context = null;
+            string computeNodeId = null;
+            string computeNodeId2 = null;
+            int originalDedicated = 3;
+            controller.RunPsTestWorkflow(
+                () => { return new string[] { string.Format("Test-RemoveMultipleComputeNodes '{0}' '{1}' '{2}' '{3}'", accountName, poolId, computeNodeId, computeNodeId2) }; },
+                () =>
+                {
+                    context = ScenarioTestHelpers.GetBatchAccountContextWithKeys(controller, accountName);
+                    originalDedicated = ScenarioTestHelpers.GetPoolCurrentDedicated(controller, context, poolId);
+                    ScenarioTestHelpers.WaitForSteadyPoolAllocation(controller, context, poolId);
+                    ScenarioTestHelpers.ResizePool(controller, context, poolId, originalDedicated + 2);
+                    ScenarioTestHelpers.WaitForSteadyPoolAllocation(controller, context, poolId);
+                    computeNodeId = ScenarioTestHelpers.GetComputeNodeId(controller, context, poolId);
+                    ScenarioTestHelpers.WaitForIdleComputeNode(controller, context, poolId, computeNodeId);
+                    computeNodeId2 = ScenarioTestHelpers.GetComputeNodeId(controller, context, poolId, 1);
+                    ScenarioTestHelpers.WaitForIdleComputeNode(controller, context, poolId, computeNodeId2);
+                },
+                null,
+                TestUtilities.GetCallingClass(),
+                TestUtilities.GetCurrentMethodName());
+        }
+
         [Fact]
         [Trait(Category.AcceptanceType, Category.CheckIn)]
         public void TestRebootComputeNodeById()

src/ResourceManager/AzureBatch/Commands.Batch.Test/ScenarioTests/ComputeNodeTests.ps1

@@ -172,7 +172,9 @@ function Test-RemoveComputeNode
     }
 
     # State transition isn't immediate
-    $computeNode = Get-AzureBatchComputeNode -PoolId $poolId -Filter "id eq '$computeNodeId'" -Select "id,state" -BatchContext $context
+    $filter = "id eq '$computeNodeId'"
+    $select = "id,state"
+    $computeNode = Get-AzureBatchComputeNode -PoolId $poolId -Filter $filter -Select $select -BatchContext $context
     $start = [DateTime]::Now
     $end = $start.AddSeconds(30)
     while ($computeNode.State -ne 'LeavingPool')
@@ -182,7 +184,36 @@ function Test-RemoveComputeNode
             throw [System.TimeoutException] "Timed out waiting for compute node to enter LeavingPool state"
         }
         Start-TestSleep 1000
-        $computeNode = Get-AzureBatchComputeNode -PoolId $poolId -Filter "id eq '$computeNodeId'" -Select "id,state" -BatchContext $context
+        $computeNode = Get-AzureBatchComputeNode -PoolId $poolId -Filter $filter -Select $select -BatchContext $context
+    }
+}
+
+<#
+.SYNOPSIS
+Tests removing multiple compute nodes from a pool
+#>
+function Test-RemoveMultipleComputeNodes
+{
+    param([string]$accountName, [string]$poolId, [string]$computeNodeId, [string]$computeNodeId2)
+
+    $context = Get-ScenarioTestContext $accountName
+
+    Remove-AzureBatchComputeNode $poolId @($computeNodeId, $computeNodeId2) -Force -BatchContext $context
+
+    # State transition isn't immediate
+    $filter = "(id eq '$computeNodeId') or (id eq '$computeNodeId2')"
+    $select = "id,state"
+    $computeNodes = Get-AzureBatchComputeNode -PoolId $poolId -Filter $filter -Select $select -BatchContext $context
+    $start = [DateTime]::Now
+    $end = $start.AddSeconds(30)
+    while ($computeNodes[0].State -ne 'LeavingPool' -and $computeNodes[1].State -ne 'LeavingPool')
+    {
+        if ([DateTime]::Now -gt $end)
+        {
+            throw [System.TimeoutException] "Timed out waiting for compute nodes to enter LeavingPool state"
+        }
+        Start-TestSleep 1000
+        $computeNodes = Get-AzureBatchComputeNode -PoolId $poolId -Filter $filter -Select $select -BatchContext $context
     }
 }
 

src/ResourceManager/AzureBatch/Commands.Batch.Test/ScenarioTests/ScenarioTestHelpers.cs

@@ -476,13 +476,13 @@ public static void TerminateJob(BatchController controller, BatchAccountContext
         /// <summary>
         /// Gets the id of a compute node in the specified pool
         /// </summary>
-        public static string GetComputeNodeId(BatchController controller, BatchAccountContext context, string poolId)
+        public static string GetComputeNodeId(BatchController controller, BatchAccountContext context, string poolId, int index = 0)
         {
             BatchClient client = new BatchClient(controller.BatchManagementClient, controller.ResourceManagementClient);
 
             ListComputeNodeOptions options = new ListComputeNodeOptions(context, poolId, null);
-
-            return client.ListComputeNodes(options).First().Id;
+            List<PSComputeNode> computeNodes = client.ListComputeNodes(options).ToList();
+            return computeNodes[index].Id;
         }
 
         /// <summary>