[DRAFT] Batch Update

src/Accounts/Accounts/AzureRmAlias/Mappings.json

@@ -356,6 +356,7 @@
     "Test-AzBatchAutoScale": "Test-AzureBatchAutoScale",
     "Get-AzBatchLocationQuotas": "Get-AzureRmBatchLocationQuotas",
     "Get-AzBatchSubtask": "Get-AzureBatchSubtask",
+    "Get-AzGetSupportedVirtualMachineSkus": "Get-AzureBatchSupportedVirtualMachineSkus",
     "Get-AzBatchTask": "Get-AzureBatchTask",
     "New-AzBatchTask": "New-AzureBatchTask",
     "Remove-AzBatchTask": "Remove-AzureBatchTask",

src/Batch/Batch.Test/Batch.Test.csproj

@@ -14,8 +14,8 @@
   </PropertyGroup>
 
   <ItemGroup>
-    <PackageReference Include="Microsoft.Azure.Batch" Version="13.0.0" />
-    <PackageReference Include="Microsoft.Azure.Management.Batch" Version="11.0.0" />
+    <PackageReference Include="Microsoft.Azure.Batch" Version="15.1.0" />
+    <PackageReference Include="Microsoft.Azure.Management.Batch" Version="14.0.0" />
     <PackageReference Include="WindowsAzure.Storage" Version="9.3.0" />
     <PackageReference Include="Microsoft.Azure.Management.Network" Version="20.6.0" />
   </ItemGroup>

src/Batch/Batch.Test/BatchAccounts/GetBatchAccountCommandTests.cs

@@ -92,7 +92,7 @@ public void GetBatchAccount_IdentityTest()
         {
             string accountName = "account01";
             string resourceGroup = "resourceGroup";
-            BatchAccount accountResource = BatchTestHelpers.CreateAccountResource(accountName, resourceGroup, identity: new BatchAccountIdentity(ResourceIdentityType.None, string.Empty, string.Empty));
+            BatchAccount accountResource = BatchTestHelpers.CreateAccountResource(accountName, resourceGroup, identity: new BatchAccountIdentity(ResourceIdentityType.None, string.Empty, string.Empty, null as IDictionary<string, UserAssignedIdentities>));
             BatchAccountContext expected = BatchAccountContext.ConvertAccountResourceToNewAccountContext(accountResource, null);
             batchClientMock.Setup(b => b.GetAccount(resourceGroup, accountName)).Returns(expected);
 

src/Batch/Batch.Test/BatchTestHelpers.cs

@@ -81,7 +81,7 @@ public static BatchAccount CreateAccountResource(
                 tags: tags == null ? null : TagsConversionHelper.CreateTagDictionary(tags, true),
                 dedicatedCoreQuotaPerVMFamilyEnforced: dedicatedCoreQuotaPerVMFamilyEnforced,
                 dedicatedCoreQuotaPerVMFamily: machineFamilyQuotas,
-                identity: identity ?? new BatchAccountIdentity(ResourceIdentityType.None));
+                identity: identity ?? new BatchAccountIdentity(ResourceIdentityType.None, null, null, null as IDictionary<string, UserAssignedIdentities>));
 
             return resource;
         }
@@ -376,18 +376,23 @@ public static RequestInterceptor ExamineRequestInterceptor<T>(Action<T> assertAc
         }
 
         /// <summary>
-        /// Builds a CloudPoolGetResponse object
+        /// Builds a CloudPoolGetResponse object using a pool ID
         /// </summary>
         public static AzureOperationResponse<ProxyModels.CloudPool, ProxyModels.PoolGetHeaders> CreateCloudPoolGetResponse(string poolId)
         {
-            var response = new AzureOperationResponse<ProxyModels.CloudPool, ProxyModels.PoolGetHeaders>();
-            response.Response = new HttpResponseMessage(HttpStatusCode.OK);
-
             ProxyModels.CloudPool pool = new ProxyModels.CloudPool();
             pool.Id = poolId;
+            return CreateCloudPoolGetResponse(pool);
+        }
 
+        /// <summary>
+        /// Builds a CloudPoolGetResponse object using a pool model
+        /// </summary>
+        public static AzureOperationResponse<ProxyModels.CloudPool, ProxyModels.PoolGetHeaders> CreateCloudPoolGetResponse(ProxyModels.CloudPool pool)
+        {
+            var response = new AzureOperationResponse<ProxyModels.CloudPool, ProxyModels.PoolGetHeaders>();
+            response.Response = new HttpResponseMessage(HttpStatusCode.OK);
             response.Body = pool;
-
             return response;
         }
 
@@ -707,20 +712,33 @@ public static AzureOperationResponse<
         /// <summary>
         /// Builds a TaskCountsGetResponse object
         /// </summary>
-        public static AzureOperationResponse<ProxyModels.TaskCounts, ProxyModels.JobGetTaskCountsHeaders> CreateTaskCountsGetResponse(
-            int active, int running, int succeeded, int failed)
+        public static AzureOperationResponse<ProxyModels.TaskCountsResult, ProxyModels.JobGetTaskCountsHeaders> CreateTaskCountsGetResponse(
+            int requiredTaskSlots, int activeTasks, int runningTasks, int succeededTasks, int failedTasks)
         {
-            var response = new AzureOperationResponse<ProxyModels.TaskCounts, ProxyModels.JobGetTaskCountsHeaders>();
+            var response = new AzureOperationResponse<ProxyModels.TaskCountsResult, ProxyModels.JobGetTaskCountsHeaders>();
             response.Response = new HttpResponseMessage(HttpStatusCode.OK);
 
+            var completedTasks = succeededTasks + failedTasks;
+
             ProxyModels.TaskCounts taskCounts = new ProxyModels.TaskCounts();
-            taskCounts.Active = active;
-            taskCounts.Running = running;
-            taskCounts.Succeeded = succeeded;
-            taskCounts.Failed = failed;
-            taskCounts.Completed = succeeded + failed;
+            taskCounts.Active = activeTasks;
+            taskCounts.Running = runningTasks;
+            taskCounts.Succeeded = succeededTasks;
+            taskCounts.Failed = failedTasks;
+            taskCounts.Completed = completedTasks;
+
+            ProxyModels.TaskSlotCounts slotCounts = new ProxyModels.TaskSlotCounts();
+            slotCounts.Active = requiredTaskSlots * activeTasks;
+            slotCounts.Running = requiredTaskSlots * runningTasks;
+            slotCounts.Succeeded = requiredTaskSlots * succeededTasks;
+            slotCounts.Failed = requiredTaskSlots * failedTasks;
+            slotCounts.Completed = requiredTaskSlots * completedTasks;
+
+            ProxyModels.TaskCountsResult result = new ProxyModels.TaskCountsResult();
+            result.TaskCounts = taskCounts;
+            result.TaskSlotCounts = slotCounts;
 
-            response.Body = taskCounts;
+            response.Body = result;
 
             return response;
         }

src/Batch/Batch.Test/Locations/GetBatchLocationQuotasCommandTests.cs

@@ -14,12 +14,14 @@
 
 using Microsoft.Azure.Commands.Batch.Models;
 using Microsoft.Azure.Management.Batch.Models;
+using Microsoft.Azure.ServiceManagement.Common.Models;
 using Microsoft.WindowsAzure.Commands.ScenarioTest;
 using Microsoft.WindowsAzure.Commands.Test.Utilities.Common;
 using Moq;
 using System.Collections.Generic;
 using System.Management.Automation;
 using Xunit;
+using Xunit.Abstractions;
 using BatchClient = Microsoft.Azure.Commands.Batch.Models.BatchClient;
 
 namespace Microsoft.Azure.Commands.Batch.Test.Subscriptions
@@ -30,9 +32,9 @@ public class GetBatchLocationQuotasCommandTests : RMTestBase
         private Mock<BatchClient> batchClientMock;
         private Mock<ICommandRuntime> commandRuntimeMock;
 
-        public GetBatchLocationQuotasCommandTests(Xunit.Abstractions.ITestOutputHelper output)
+        public GetBatchLocationQuotasCommandTests(ITestOutputHelper output)
         {
-            ServiceManagement.Common.Models.XunitTracingInterceptor.AddToContext(new ServiceManagement.Common.Models.XunitTracingInterceptor(output));
+            XunitTracingInterceptor.AddToContext(new XunitTracingInterceptor(output));
             batchClientMock = new Mock<BatchClient>();
             commandRuntimeMock = new Mock<ICommandRuntime>();
             cmdlet = new GetBatchLocationQuotaCommand()
@@ -46,8 +48,6 @@ public GetBatchLocationQuotasCommandTests(Xunit.Abstractions.ITestOutputHelper o
         [Trait(Category.AcceptanceType, Category.CheckIn)]
         public void GetBatchLocationQuotasTest()
         {
-            List<PSBatchLocationQuotas> pipelineOutput = new List<PSBatchLocationQuotas>();
-
             // Return a pre-built object when the command is issued.
             string location = "westus";
             PSBatchLocationQuotas quotas = new PSBatchLocationQuotas(location, new BatchLocationQuota(accountQuota: 5));

src/Batch/Batch.Test/Locations/GetSupportedVirtualMachineSkusTests.cs

@@ -0,0 +1,110 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+
+using Microsoft.Azure.Commands.Batch;
+using Microsoft.Azure.Commands.Batch.Models;
+using Microsoft.Azure.Management.Batch.Models;
+using Microsoft.Azure.ServiceManagement.Common.Models;
+using Microsoft.WindowsAzure.Commands.ScenarioTest;
+using Microsoft.WindowsAzure.Commands.Test.Utilities.Common;
+using Moq;
+using System.Collections.Generic;
+using System.Linq;
+using System.Management.Automation;
+using Xunit;
+using Xunit.Abstractions;
+using BatchClient = Microsoft.Azure.Commands.Batch.Models.BatchClient;
+
+namespace Microsoft.Azure.Commands.Batch.Test.Subscriptions
+{
+    public class GetSupportedVirtualMachineSkusTests : RMTestBase
+    {
+        private GetSupportedVirtualMachineSkusCommand cmdlet;
+        private Mock<BatchClient> batchClientMock;
+        private Mock<ICommandRuntime> commandRuntimeMock;
+
+        public GetSupportedVirtualMachineSkusTests(ITestOutputHelper output)
+        {
+            XunitTracingInterceptor.AddToContext(new XunitTracingInterceptor(output));
+            batchClientMock = new Mock<BatchClient>();
+            commandRuntimeMock = new Mock<ICommandRuntime>();
+            cmdlet = new GetSupportedVirtualMachineSkusCommand()
+            {
+                CommandRuntime = commandRuntimeMock.Object,
+                BatchClient = batchClientMock.Object
+            };
+        }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void TestGetSupportedSkus()
+        {
+            string location = "westus";
+
+            var skus = CreateSkus();
+            batchClientMock.Setup(client => client.GetSupportedVirtualMachineSkus(location, default, default)).Returns(skus);
+
+            cmdlet.Location = location;
+            cmdlet.ExecuteCmdlet();
+
+            commandRuntimeMock.Verify(r => r.WriteObject(skus, true), Times.Once());
+        }
+
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void TestGetSupportedSkusWithODataParameters()
+        {
+            string location = "westus";
+
+            var skus = CreateSkus();
+            batchClientMock.Setup(client => client.GetSupportedVirtualMachineSkus(location, 3, "testfilter")).Returns(skus);
+
+            cmdlet.Location = location;
+            cmdlet.MaxResults = 3;
+            cmdlet.Filter = "testfilter";
+            cmdlet.ExecuteCmdlet();
+
+            commandRuntimeMock.Verify(r => r.WriteObject(skus, true), Times.Once());
+        }
+
+        private List<PSSupportedSku> CreateSkus()
+        {
+            List<PSSupportedSku> skus = new List<PSSupportedSku>()
+            {
+                new PSSupportedSku("testsku1", "testfamily", new List<PSSkuCapability>()
+                {
+                    new PSSkuCapability("cap1", "val1"),
+                    new PSSkuCapability("cap2", "val2")
+                }),
+                new PSSupportedSku("testsku2", "testfamily", new List<PSSkuCapability>()
+                {
+                    new PSSkuCapability("cap1", "val1"),
+                }),
+                new PSSupportedSku("testsku3", "testfamily", new List<PSSkuCapability>()
+                {
+                    new PSSkuCapability("cap1", "val1"),
+                    new PSSkuCapability("cap2", "val2"),
+                    new PSSkuCapability("cap2", "val2")
+                }),
+                new PSSupportedSku("testsku4", "testfamily", new List<PSSkuCapability>()
+                {
+                    new PSSkuCapability("cap1", "val1"),
+                }),
+            };
+
+            return skus;
+        }
+    }
+}

src/Batch/Batch.Test/Pools/GetBatchPoolCommandTests.cs

@@ -57,8 +57,13 @@ public void GetBatchPoolTest()
             cmdlet.Id = "testPool";
             cmdlet.Filter = null;
 
+            // Pool returned in the response
+            ProxyModels.CloudPool pool = new ProxyModels.CloudPool();
+            pool.Id = cmdlet.Id;
+            pool.TaskSlotsPerNode = 16;
+
             // Build a CloudPool instead of querying the service on a Get CloudPool call
-            AzureOperationResponse<ProxyModels.CloudPool, ProxyModels.PoolGetHeaders> response = BatchTestHelpers.CreateCloudPoolGetResponse(cmdlet.Id);
+            AzureOperationResponse<ProxyModels.CloudPool, ProxyModels.PoolGetHeaders> response = BatchTestHelpers.CreateCloudPoolGetResponse(pool);
             RequestInterceptor interceptor = BatchTestHelpers.CreateFakeServiceResponseInterceptor<
                 ProxyModels.PoolGetOptions,
                 AzureOperationResponse<ProxyModels.CloudPool, ProxyModels.PoolGetHeaders>>(response);
@@ -74,6 +79,7 @@ public void GetBatchPoolTest()
             // Verify that the cmdlet wrote the pool returned from the OM to the pipeline
             Assert.Single(pipeline);
             Assert.Equal(cmdlet.Id, pipeline[0].Id);
+            Assert.Equal(16, pipeline[0].TaskSlotsPerNode);
         }
 
         [Fact]

src/Batch/Batch.Test/Pools/NewBatchPoolCommandTests.cs

@@ -93,7 +93,7 @@ public void NewBatchPoolParametersGetPassedToRequestTest()
             cmdlet.CloudServiceConfiguration = new PSCloudServiceConfiguration("4", "*");
             cmdlet.DisplayName = "display name";
             cmdlet.InterComputeNodeCommunicationEnabled = true;
-            cmdlet.MaxTasksPerComputeNode = 4;
+            cmdlet.TaskSlotsPerNode = 4;
             cmdlet.Metadata = new Dictionary<string, string>();
             cmdlet.Metadata.Add("meta1", "value1");
             cmdlet.ResizeTimeout = TimeSpan.FromMinutes(20);
@@ -103,7 +103,10 @@ public void NewBatchPoolParametersGetPassedToRequestTest()
             cmdlet.TaskSchedulingPolicy = new PSTaskSchedulingPolicy(Azure.Batch.Common.ComputeNodeFillType.Spread);
             cmdlet.VirtualMachineConfiguration = new PSVirtualMachineConfiguration(new PSImageReference("offer", "publisher", "sku"), "node agent");
             cmdlet.VirtualMachineSize = "small";
-            cmdlet.MountConfiguration = new[] { new PSMountConfiguration(new PSAzureBlobFileSystemConfiguration("foo", "bar", "baz", AzureStorageAuthenticationKey.FromAccountKey("abc"))) };
+            cmdlet.MountConfiguration = new[] {
+                new PSMountConfiguration(new PSAzureBlobFileSystemConfiguration("foo", "bar", "baz", AzureStorageAuthenticationKey.FromAccountKey("abc"))),
+                new PSMountConfiguration(new PSAzureBlobFileSystemConfiguration("foo2", "bar2", "baz2", new PSComputeNodeIdentityReference(new Azure.Batch.ComputeNodeIdentityReference { ResourceId = "fake-identity"})))
+            };
 
             PoolAddParameter requestParameters = null;
 
@@ -130,7 +133,7 @@ public void NewBatchPoolParametersGetPassedToRequestTest()
             Assert.Equal(cmdlet.CloudServiceConfiguration.OSVersion, requestParameters.CloudServiceConfiguration.OsVersion);
             Assert.Equal(cmdlet.DisplayName, requestParameters.DisplayName);
             Assert.Equal(cmdlet.InterComputeNodeCommunicationEnabled, requestParameters.EnableInterNodeCommunication);
-            Assert.Equal(cmdlet.MaxTasksPerComputeNode, requestParameters.MaxTasksPerNode);
+            Assert.Equal(cmdlet.TaskSlotsPerNode, requestParameters.TaskSlotsPerNode);
             Assert.Equal(cmdlet.Metadata.Count, requestParameters.Metadata.Count);
             Assert.Equal(cmdlet.Metadata["meta1"], requestParameters.Metadata[0].Value);
             Assert.Equal(cmdlet.ResizeTimeout, requestParameters.ResizeTimeout);
@@ -147,6 +150,10 @@ public void NewBatchPoolParametersGetPassedToRequestTest()
             Assert.Equal(cmdlet.MountConfiguration[0].AzureBlobFileSystemConfiguration.AccountKey, requestParameters.MountConfiguration[0].AzureBlobFileSystemConfiguration.AccountKey);
             Assert.Equal(cmdlet.MountConfiguration[0].AzureBlobFileSystemConfiguration.ContainerName, requestParameters.MountConfiguration[0].AzureBlobFileSystemConfiguration.ContainerName);
             Assert.Equal(cmdlet.MountConfiguration[0].AzureBlobFileSystemConfiguration.RelativeMountPath, requestParameters.MountConfiguration[0].AzureBlobFileSystemConfiguration.RelativeMountPath);
+            Assert.Equal(cmdlet.MountConfiguration[1].AzureBlobFileSystemConfiguration.AccountName, requestParameters.MountConfiguration[1].AzureBlobFileSystemConfiguration.AccountName);
+            Assert.Equal(cmdlet.MountConfiguration[1].AzureBlobFileSystemConfiguration.IdentityReference.ResourceId, requestParameters.MountConfiguration[1].AzureBlobFileSystemConfiguration.IdentityReference.ResourceId);
+            Assert.Equal(cmdlet.MountConfiguration[1].AzureBlobFileSystemConfiguration.ContainerName, requestParameters.MountConfiguration[1].AzureBlobFileSystemConfiguration.ContainerName);
+            Assert.Equal(cmdlet.MountConfiguration[1].AzureBlobFileSystemConfiguration.RelativeMountPath, requestParameters.MountConfiguration[1].AzureBlobFileSystemConfiguration.RelativeMountPath);
         }
 
         [Fact]

src/Batch/Batch.Test/ScenarioTests/CertificateTests.ps1

@@ -62,8 +62,8 @@ function Test-TestCancelCertificateDelete
     Get-AzBatchCertificate $thumbprintAlgorithm $thumbprint -BatchContext $context | Stop-AzBatchCertificateDeletion -BatchContext $context
 
     # Verify the cert went back to the active state
-    $filter = "state eq 'active'"
+    $filter = "state eq 'active'";
     $cert = Get-AzBatchCertificate -Filter $filter -BatchContext $context
 
-    Assert-AreEqual $thumbprint $cert.Thumbprint.ToLowerInvariant()
+    Assert-True { $cert.Thumbprint.ToLowerInvariant() -Contains $thumbprint }
 }

src/Batch/Batch.Test/ScenarioTests/JobTests.ps1

@@ -114,7 +114,7 @@ function Test-JobWithTaskDependencies
 
         $poolSpec = New-Object Microsoft.Azure.Commands.Batch.Models.PSPoolSpecification
         $poolSpec.TargetDedicated = $targetDedicated = 3
-        $poolSpec.VirtualMachineSize = $vmSize = "small"
+        $poolSpec.VirtualMachineSize = $vmSize = "standard_d1_v2"
         $poolSpec.CloudServiceConfiguration = $paasConfiguration
         $autoPoolSpec = New-Object Microsoft.Azure.Commands.Batch.Models.PSAutoPoolSpecification
         $autoPoolSpec.PoolSpecification = $poolSpec
@@ -162,8 +162,8 @@ function IfJobSetsAutoFailure-ItCompletesWhenAnyTaskFails
     $paasConfiguration = New-Object Microsoft.Azure.Commands.Batch.Models.PSCloudServiceConfiguration -ArgumentList @($osFamily, $targetOSVersion)
 
     $poolSpec = New-Object Microsoft.Azure.Commands.Batch.Models.PSPoolSpecification
-    $poolSpec.TargetDedicatedComputeNodes = $targetDedicated = 3
-    $poolSpec.VirtualMachineSize = $vmSize = "small"
+    $poolSpec.TargetDedicatedComputeNodes = $targetDedicated = 1
+    $poolSpec.VirtualMachineSize = $vmSize = "standard_d1_v2"
     $poolSpec.CloudServiceConfiguration = $paasConfiguration
     $autoPoolSpec = New-Object Microsoft.Azure.Commands.Batch.Models.PSAutoPoolSpecification
     $autoPoolSpec.PoolSpecification = $poolSpec