[Az.Batch] Adds Get-AzBatchSupportedVirtualMachineSku, Get-AzBatchTaskSlotCount, and Get-AzBatchComputeNodeExtension commands

src/Accounts/Accounts/AzureRmAlias/Mappings.json

@@ -356,6 +356,7 @@
     "Test-AzBatchAutoScale": "Test-AzureBatchAutoScale",
     "Get-AzBatchLocationQuotas": "Get-AzureRmBatchLocationQuotas",
     "Get-AzBatchSubtask": "Get-AzureBatchSubtask",
+    "Get-AzBatchSupportedVirtualMachineSku": "Get-AzureBatchSupportedVirtualMachineSku",
     "Get-AzBatchTask": "Get-AzureBatchTask",
     "New-AzBatchTask": "New-AzureBatchTask",
     "Remove-AzBatchTask": "Remove-AzureBatchTask",
@@ -365,6 +366,7 @@
     "Set-AzBatchTask": "Set-AzureBatchTask",
     "Stop-AzBatchTask": "Stop-AzureBatchTask",
     "Get-AzBatchComputeNode": "Get-AzureBatchComputeNode",
+    "Get-AzBatchComputeNodeExtension": "Get-AzureBatchComputeNodeExtension",
     "Get-AzBatchJobSchedule": "Get-AzureBatchJobSchedule",
     "New-AzBatchJobSchedule": "New-AzureBatchJobSchedule",
     "Remove-AzBatchJobSchedule": "Remove-AzureBatchJobSchedule",

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

@@ -13,8 +13,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.3.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="21.0.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);
 
             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;
         }
 
@@ -552,6 +557,50 @@ public static RequestInterceptor ExamineRequestInterceptor<T>(Action<T> assertAc
             return response;
         }
 
+        /// <summary>
+        /// Builds a ComputeNodeExtensionGetResponse object
+        /// </summary>
+        public static AzureOperationResponse<ProxyModels.NodeVMExtension, ProxyModels.ComputeNodeExtensionGetHeaders> CreateComputeNodeExtensionGetResponse(Azure.Batch.VMExtension extension)
+        {
+            var response = new AzureOperationResponse<ProxyModels.NodeVMExtension, ProxyModels.ComputeNodeExtensionGetHeaders>();
+            response.Response = new HttpResponseMessage(HttpStatusCode.OK);
+
+            ProxyModels.NodeVMExtension proxyExtension = CreateProxyExtension(extension);
+            response.Body = proxyExtension;
+
+            return response;
+        }
+
+        /// <summary>
+        /// Builds a ComputeNodeExtensionListResponse object
+        /// </summary>
+        public static AzureOperationResponse<IPage<ProxyModels.NodeVMExtension>, ProxyModels.ComputeNodeExtensionListHeaders> CreateComputeNodeExtensionListResponse(IEnumerable<Azure.Batch.VMExtension> extensions)
+        {
+            var response = new AzureOperationResponse<IPage<ProxyModels.NodeVMExtension>, ProxyModels.ComputeNodeExtensionListHeaders>();
+            response.Response = new HttpResponseMessage(HttpStatusCode.OK);
+
+            List<ProxyModels.NodeVMExtension> proxyExtensions = new List<ProxyModels.NodeVMExtension>();
+
+            foreach (Azure.Batch.VMExtension extension in extensions)
+            {
+                ProxyModels.NodeVMExtension proxyExtension = CreateProxyExtension(extension);
+                proxyExtensions.Add(proxyExtension);
+            }
+
+            response.Body = new MockPagedEnumerable<ProxyModels.NodeVMExtension>(proxyExtensions);
+
+            return response;
+        }
+
+        private static ProxyModels.NodeVMExtension CreateProxyExtension(Azure.Batch.VMExtension extension)
+        {
+            ProxyModels.NodeVMExtension proxyExtension = new ProxyModels.NodeVMExtension();
+            proxyExtension.InstanceView = new ProxyModels.VMExtensionInstanceView();
+            proxyExtension.VmExtension = new ProxyModels.VMExtension(extension.Name, extension.Publisher, extension.Type);
+            proxyExtension.ProvisioningState = ProvisioningState.Succeeded.ToString();
+            return proxyExtension;
+        }
+
         /// <summary>
         /// Builds a CloudJobScheduleGetResponse object
         /// </summary>
@@ -707,20 +756,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 slotCount = new ProxyModels.TaskSlotCounts();
+            slotCount.Active = requiredTaskSlots * activeTasks;
+            slotCount.Running = requiredTaskSlots * runningTasks;
+            slotCount.Succeeded = requiredTaskSlots * succeededTasks;
+            slotCount.Failed = requiredTaskSlots * failedTasks;
+            slotCount.Completed = requiredTaskSlots * completedTasks;
+
+            ProxyModels.TaskCountsResult result = new ProxyModels.TaskCountsResult();
+            result.TaskCounts = taskCounts;
+            result.TaskSlotCounts = slotCount;
 
-            response.Body = taskCounts;
+            response.Body = result;
 
             return response;
         }

src/Batch/Batch.Test/ComputeNodeExtensions/GetBatchComputeNodeExtensionCommandTests.cs

@@ -0,0 +1,218 @@
+// ----------------------------------------------------------------------------------
+//
+// 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.Batch;
+using Microsoft.Azure.Batch.Protocol;
+using Microsoft.Azure.Commands.Batch.Models;
+using Microsoft.Rest.Azure;
+using Microsoft.WindowsAzure.Commands.ScenarioTest;
+using Moq;
+using System;
+using System.Collections.Generic;
+using System.Diagnostics;
+using System.Linq;
+using System.Management.Automation;
+using System.Threading.Tasks;
+using Xunit;
+using BatchClient = Microsoft.Azure.Commands.Batch.Models.BatchClient;
+using ProxyModels = Microsoft.Azure.Batch.Protocol.Models;
+
+namespace Microsoft.Azure.Commands.Batch.Test.ComputeNodeExtensions
+{
+    public class GetBatchComputeNodeExtensionCommandTests : WindowsAzure.Commands.Test.Utilities.Common.RMTestBase
+    {
+        private GetBatchComputeNodeExtensionCommand cmdlet;
+        private Mock<BatchClient> batchClientMock;
+        private Mock<ICommandRuntime> commandRuntimeMock;
+
+        public GetBatchComputeNodeExtensionCommandTests(Xunit.Abstractions.ITestOutputHelper output)
+        {
+            ServiceManagement.Common.Models.XunitTracingInterceptor.AddToContext(new ServiceManagement.Common.Models.XunitTracingInterceptor(output));
+            batchClientMock = new Mock<BatchClient>();
+            commandRuntimeMock = new Mock<ICommandRuntime>();
+            cmdlet = new GetBatchComputeNodeExtensionCommand()
+            {
+                CommandRuntime = commandRuntimeMock.Object,
+                BatchClient = batchClientMock.Object,
+            };
+        }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void GetBatchComputeNodeExtensionTest()
+        {
+            string extensionName = "testExtension";
+            string publisher = "testPublisher";
+            string type = "testType";
+
+            // Setup cmdlet to get a compute node by id
+            BatchAccountContext context = BatchTestHelpers.CreateBatchContextWithKeys();
+            cmdlet.BatchContext = context;
+            cmdlet.PoolId = "testPool";
+            cmdlet.ComputeNodeId = "testComputeNode";
+            cmdlet.Name = extensionName;
+
+            VMExtension extension = new VMExtension(extensionName, publisher, type);
+
+            // Build an extension instead of querying the service on a Get ComputeNodeExtension call
+            AzureOperationResponse<ProxyModels.NodeVMExtension, ProxyModels.ComputeNodeExtensionGetHeaders> response = BatchTestHelpers.CreateComputeNodeExtensionGetResponse(extension);
+            RequestInterceptor interceptor = BatchTestHelpers.CreateFakeServiceResponseInterceptor<
+                ProxyModels.ComputeNodeExtensionGetOptions,
+                AzureOperationResponse<ProxyModels.NodeVMExtension, ProxyModels.ComputeNodeExtensionGetHeaders>>(response);
+
+            cmdlet.AdditionalBehaviors = new List<BatchClientBehavior>() { interceptor };
+
+            // Setup the cmdlet to write pipeline output to a list that can be examined later
+            List<PSNodeVMExtension> pipeline = new List<PSNodeVMExtension>();
+            commandRuntimeMock.Setup(r => r.WriteObject(It.IsAny<PSNodeVMExtension>())).Callback<object>(c => pipeline.Add((PSNodeVMExtension)c));
+
+            cmdlet.ExecuteCmdlet();
+
+            // Verify that the cmdlet wrote the compute node returned from the OM to the pipeline
+            Assert.Single(pipeline);
+
+            PSVMExtension pipelineExtension = pipeline[0].VmExtension;
+            Assert.NotNull(pipelineExtension);
+            Assert.Equal("testExtension", pipelineExtension.Name);
+            Assert.Equal("testPublisher", pipelineExtension.Publisher);
+            Assert.Equal("testType", pipelineExtension.Type);
+        }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void GetBatchComputeNodeExtensionODataTest()
+        {
+            string extensionName = "testExtension";
+            string publisher = "testPublisher";
+            string type = "testType";
+
+            // Setup cmdlet to get a compute node by id
+            BatchAccountContext context = BatchTestHelpers.CreateBatchContextWithKeys();
+            cmdlet.BatchContext = context;
+            cmdlet.PoolId = "testPool";
+            cmdlet.ComputeNodeId = "testComputeNode";
+            cmdlet.Name = extensionName;
+            cmdlet.Select = "ExtensionName,Publisher";
+
+            VMExtension extension = new VMExtension(extensionName, publisher, type);
+
+            // Fetch the OData clauses off the request. The OData clauses are applied after user provided RequestInterceptors, so a ResponseInterceptor is used.
+            AzureOperationResponse<ProxyModels.NodeVMExtension, ProxyModels.ComputeNodeExtensionGetHeaders> getResponse = BatchTestHelpers.CreateComputeNodeExtensionGetResponse(extension);
+            RequestInterceptor requestInterceptor = BatchTestHelpers.CreateFakeServiceResponseInterceptor<
+                ProxyModels.ComputeNodeExtensionGetOptions,
+                AzureOperationResponse<ProxyModels.NodeVMExtension, ProxyModels.ComputeNodeExtensionGetHeaders>>(getResponse);
+
+            string requestSelect = null;
+
+            ResponseInterceptor responseInterceptor = new ResponseInterceptor((response, request) =>
+            {
+                ProxyModels.ComputeNodeExtensionGetOptions options = (ProxyModels.ComputeNodeExtensionGetOptions)request.Options;
+                requestSelect = options.Select;
+
+                return Task.FromResult(response);
+            });
+
+            cmdlet.AdditionalBehaviors = new List<BatchClientBehavior>() { requestInterceptor, responseInterceptor };
+
+            cmdlet.ExecuteCmdlet();
+
+            Assert.Equal(cmdlet.Select, requestSelect);
+        }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void ListBatchComputeNodeExtensionsTest()
+        {
+            // Setup cmdlet to get a compute node by id
+            BatchAccountContext context = BatchTestHelpers.CreateBatchContextWithKeys();
+            cmdlet.BatchContext = context;
+            cmdlet.PoolId = "testPool";
+            cmdlet.ComputeNodeId = "testComputeNode";
+
+            int count = 5;
+
+            // Build an extension instead of querying the service on a Get ComputeNodeExtension call
+            AzureOperationResponse<IPage<ProxyModels.NodeVMExtension>, ProxyModels.ComputeNodeExtensionListHeaders> response = BatchTestHelpers.CreateComputeNodeExtensionListResponse(CreateTestExtensions(count));
+            RequestInterceptor interceptor = BatchTestHelpers.CreateFakeServiceResponseInterceptor<
+                ProxyModels.ComputeNodeExtensionListOptions,
+                AzureOperationResponse<IPage<ProxyModels.NodeVMExtension>, ProxyModels.ComputeNodeExtensionListHeaders>>(response);
+
+            cmdlet.AdditionalBehaviors = new List<BatchClientBehavior>() { interceptor };
+
+            // Setup the cmdlet to write pipeline output to a list that can be examined later
+            List<PSNodeVMExtension> pipeline = new List<PSNodeVMExtension>();
+            commandRuntimeMock.Setup(r => r.WriteObject(It.IsAny<PSNodeVMExtension>())).Callback<object>(c => pipeline.Add((PSNodeVMExtension)c));
+
+            cmdlet.ExecuteCmdlet();
+
+            // Verify that the cmdlet wrote the compute node returned from the OM to the pipeline
+            Assert.Equal(count, pipeline.Count);
+            for (int i = 1; i <= count; i++)
+            {
+                PSVMExtension extension = pipeline[i-1].VmExtension;
+                Assert.Equal($"testExtension{i}", extension.Name);
+                Assert.Equal($"testPublisher{i}", extension.Publisher);
+                Assert.Equal($"testType{i}", extension.Type);
+            }
+        }
+
+        private IEnumerable<VMExtension> CreateTestExtensions(int count)
+        {
+            for (int i = 1; i <= count; i++)
+            {
+                yield return new VMExtension($"testExtension{i}", $"testPublisher{i}", $"testType{i}");
+            }
+        }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void ListBatchComputeNodeExtensionsWithMaxCountTest()
+        {
+            int maxCount = 3;
+
+            // Setup cmdlet to get a compute node by id
+            BatchAccountContext context = BatchTestHelpers.CreateBatchContextWithKeys();
+            cmdlet.BatchContext = context;
+            cmdlet.PoolId = "testPool";
+            cmdlet.ComputeNodeId = "testComputeNode";
+            cmdlet.MaxCount = maxCount;
+
+            int count = 5;
+
+            // Build an extension instead of querying the service on a Get ComputeNodeExtension call
+            AzureOperationResponse<IPage<ProxyModels.NodeVMExtension>, ProxyModels.ComputeNodeExtensionListHeaders> response = BatchTestHelpers.CreateComputeNodeExtensionListResponse(CreateTestExtensions(count));
+            RequestInterceptor interceptor = BatchTestHelpers.CreateFakeServiceResponseInterceptor<
+                ProxyModels.ComputeNodeExtensionListOptions,
+                AzureOperationResponse<IPage<ProxyModels.NodeVMExtension>, ProxyModels.ComputeNodeExtensionListHeaders>>(response);
+
+            cmdlet.AdditionalBehaviors = new List<BatchClientBehavior>() { interceptor };
+
+            // Setup the cmdlet to write pipeline output to a list that can be examined later
+            List<PSNodeVMExtension> pipeline = new List<PSNodeVMExtension>();
+            commandRuntimeMock.Setup(r => r.WriteObject(It.IsAny<PSNodeVMExtension>())).Callback<object>(c => pipeline.Add((PSNodeVMExtension)c));
+
+            cmdlet.ExecuteCmdlet();
+
+            // Verify that the cmdlet wrote the compute node returned from the OM to the pipeline
+            Assert.Equal(maxCount, pipeline.Count);
+            for (int i = 1; i <= maxCount; i++)
+            {
+                PSVMExtension extension = pipeline[i - 1].VmExtension;
+                Assert.Equal($"testExtension{i}", extension.Name);
+                Assert.Equal($"testPublisher{i}", extension.Publisher);
+                Assert.Equal($"testType{i}", extension.Type);
+            }
+        }
+    }
+}