Implement Get-AzureBatchPoolNodeCounts and Start-AzureBatchComputeNodeServiceLogUpload cmdlets

src/ResourceManager/AzureBatch/AzureRM.Batch.Netcore.psd1

@@ -107,8 +107,9 @@ CmdletsToExport = 'Remove-AzureRmBatchAccount', 'Get-AzureRmBatchAccount',
     'New-AzureBatchComputeNodeUser', 'Remove-AzureBatchComputeNodeUser', 
     'Enable-AzureBatchTask', 'Set-AzureBatchTask', 'Stop-AzureBatchTask', 
     'Get-AzureBatchComputeNode', 'Get-AzureBatchJobSchedule', 
-    'New-AzureBatchJobSchedule', 'Remove-AzureBatchJobSchedule', 
-    'Get-AzureBatchTaskCounts'
+    'New-AzureBatchJobSchedule', 'Remove-AzureBatchJobSchedule',
+    'Get-AzureBatchTaskCounts', 'Get-AzureBatchPoolNodeCounts',
+    'Start-AzureBatchComputeNodeServiceLogUpload'
 
 # Variables to export from this module
 # VariablesToExport = @()

src/ResourceManager/AzureBatch/AzureRM.Batch.psd1

@@ -111,7 +111,8 @@ CmdletsToExport = 'Remove-AzureRmBatchAccount', 'Get-AzureRmBatchAccount',
                'Enable-AzureBatchTask', 'Set-AzureBatchTask', 'Stop-AzureBatchTask', 
                'Get-AzureBatchComputeNode', 'Get-AzureBatchJobSchedule', 
                'New-AzureBatchJobSchedule', 'Remove-AzureBatchJobSchedule', 
-               'Get-AzureBatchTaskCounts'
+               'Get-AzureBatchTaskCounts', 'Get-AzureBatchPoolNodeCounts',
+               'Start-AzureBatchComputeNodeServiceLogUpload'
 
 # Variables to export from this module
 # VariablesToExport = @()

src/ResourceManager/AzureBatch/BatchModelGenerator/Program.cs

@@ -84,6 +84,7 @@ public class Program
             {"Microsoft.Azure.Batch.NetworkConfiguration", "PSNetworkConfiguration"},
             {"Microsoft.Azure.Batch.NetworkSecurityGroupRule", "PSNetworkSecurityGroupRule"},
             {"Microsoft.Azure.Batch.NodeAgentSku", "PSNodeAgentSku"},
+            {"Microsoft.Azure.Batch.NodeCounts", "PSNodeCounts"},
             {"Microsoft.Azure.Batch.NodeFile", "PSNodeFile"},
             {"Microsoft.Azure.Batch.OSDisk", "PSOSDisk"},
             {"Microsoft.Azure.Batch.OutputFile", "PSOutputFile"},
@@ -92,6 +93,7 @@ public class Program
             {"Microsoft.Azure.Batch.OutputFileBlobContainerDestination", "PSOutputFileBlobContainerDestination"},
             {"Microsoft.Azure.Batch.PoolEndpointConfiguration", "PSPoolEndpointConfiguration"},
             {"Microsoft.Azure.Batch.PoolInformation", "PSPoolInformation"},
+            {"Microsoft.Azure.Batch.PoolNodeCounts", "PSPoolNodeCounts"},
             {"Microsoft.Azure.Batch.PoolSpecification", "PSPoolSpecification"},
             {"Microsoft.Azure.Batch.PoolStatistics", "PSPoolStatistics"},
             {"Microsoft.Azure.Batch.PoolUsageMetrics", "PSPoolUsageMetrics"},
@@ -114,6 +116,7 @@ public class Program
             {"Microsoft.Azure.Batch.TaskFailureInformation", "PSTaskFailureInformation"},
             {"Microsoft.Azure.Batch.TaskSchedulingPolicy", "PSTaskSchedulingPolicy"},
             {"Microsoft.Azure.Batch.TaskStatistics", "PSTaskStatistics"},
+            {"Microsoft.Azure.Batch.UploadBatchServiceLogsResult", "PSStartComputeNodeServiceLogUploadResult"},
             {"Microsoft.Azure.Batch.UsageStatistics", "PSUsageStatistics"},
             {"Microsoft.Azure.Batch.UserAccount", "PSUserAccount"},
             {"Microsoft.Azure.Batch.UserIdentity", "PSUserIdentity"},

src/ResourceManager/AzureBatch/ChangeLog.md

@@ -18,6 +18,8 @@
         - Additional information about change #1
 -->
 ## Current Release
+* Release new cmdlet Get-AzureBatchPoolNodeCounts
+* Release new cmdlet Start-AzureBatchComputeNodeServiceLogUpload
 
 ## Version 4.0.7
 * Set minimum dependency of module to PowerShell 5.0

src/ResourceManager/AzureBatch/Commands.Batch.Test/BatchTestHelpers.cs

@@ -712,6 +712,17 @@ public static AzureOperationResponse<
             return response;
         }
 
+        public static AzureOperationResponse<IPage<ProxyModels.PoolNodeCounts>, ProxyModels.AccountListPoolNodeCountsHeaders> CreatePoolNodeCountsGetResponse(
+            IEnumerable<ProxyModels.PoolNodeCounts> poolNodeCounts)
+        {
+            var response = new AzureOperationResponse<IPage<ProxyModels.PoolNodeCounts>, ProxyModels.AccountListPoolNodeCountsHeaders>();
+            response.Response = new HttpResponseMessage(HttpStatusCode.OK);
+
+            response.Body = new MockPagedEnumerable<ProxyModels.PoolNodeCounts>(poolNodeCounts);
+
+            return response;
+        }
+
         /// <summary>
         /// Builds a CloudTaskListResponse object
         /// </summary>
@@ -980,6 +991,23 @@ public static CloudTask CreateFakeBoundTask(BatchAccountContext context)
             return context.BatchOMClient.JobOperations.GetTask("jobId", taskId, additionalBehaviors: new BatchClientBehavior[] { interceptor });
         }
 
+        /// <summary>
+        /// Builds a ComputeNodeGetResponse object
+        /// </summary>
+        public static AzureOperationResponse<ProxyModels.UploadBatchServiceLogsResult, ProxyModels.ComputeNodeUploadBatchServiceLogsHeaders> CreateComputeNodeServiceLogsAddResponse(int numberOfFilesUploaded, string virtualDirectoryName)
+        {
+            var response = new AzureOperationResponse<ProxyModels.UploadBatchServiceLogsResult, ProxyModels.ComputeNodeUploadBatchServiceLogsHeaders>();
+            response.Response = new HttpResponseMessage(HttpStatusCode.OK);
+
+            ProxyModels.UploadBatchServiceLogsResult result = new ProxyModels.UploadBatchServiceLogsResult();
+            result.NumberOfFilesUploaded = numberOfFilesUploaded;
+            result.VirtualDirectoryName = virtualDirectoryName;
+
+            response.Body = result;
+
+            return response;
+        }
+
         /// <summary>
         /// Uses Reflection to set a property value on an object. Can be used to bypass restricted set accessors.
         /// </summary>

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

@@ -45,7 +45,7 @@
       <HintPath>..\..\..\packages\Hyak.Common.1.0.3\lib\portable-net403+win+wpa81\Hyak.Common.dll</HintPath>
     </Reference>
     <Reference Include="Microsoft.Azure.Batch, Version=8.0.0.0, Culture=neutral, PublicKeyToken=31bf3856ad364e35, processorArchitecture=MSIL">
-      <HintPath>..\..\..\packages\Azure.Batch.8.0.0\lib\net452\Microsoft.Azure.Batch.dll</HintPath>
+      <HintPath>..\..\..\packages\Azure.Batch.8.1.1\lib\net452\Microsoft.Azure.Batch.dll</HintPath>
     </Reference>
     <Reference Include="Microsoft.Azure.Common">
       <HintPath>..\..\..\packages\Microsoft.Azure.Common.2.1.0\lib\net45\Microsoft.Azure.Common.dll</HintPath>
@@ -93,14 +93,14 @@
       <SpecificVersion>False</SpecificVersion>
       <HintPath>..\..\..\packages\Microsoft.IdentityModel.Clients.ActiveDirectory.2.28.3\lib\net45\Microsoft.IdentityModel.Clients.ActiveDirectory.dll</HintPath>
     </Reference>
+    <Reference Include="Microsoft.Rest.ClientRuntime.Azure.Authentication, Version=2.0.0.0, Culture=neutral, PublicKeyToken=31bf3856ad364e35, processorArchitecture=MSIL"> 
+      <HintPath>..\..\..\packages\Microsoft.Rest.ClientRuntime.Azure.Authentication.2.2.9-preview\lib\net45\Microsoft.Rest.ClientRuntime.Azure.Authentication.dll</HintPath> 
+      <Private>True</Private> 
+    </Reference>
     <Reference Include="Microsoft.IdentityModel.Clients.ActiveDirectory.WindowsForms, Version=2.28.3.860, Culture=neutral, PublicKeyToken=31bf3856ad364e35, processorArchitecture=MSIL">
       <SpecificVersion>False</SpecificVersion>
       <HintPath>..\..\..\packages\Microsoft.IdentityModel.Clients.ActiveDirectory.2.28.3\lib\net45\Microsoft.IdentityModel.Clients.ActiveDirectory.WindowsForms.dll</HintPath>
     </Reference>
-    <Reference Include="Microsoft.Rest.ClientRuntime.Azure.Authentication, Version=2.0.0.0, Culture=neutral, PublicKeyToken=31bf3856ad364e35, processorArchitecture=MSIL">
-      <HintPath>..\..\..\packages\Microsoft.Rest.ClientRuntime.Azure.Authentication.2.2.9-preview\lib\net45\Microsoft.Rest.ClientRuntime.Azure.Authentication.dll</HintPath>
-      <Private>True</Private>
-    </Reference>
     <Reference Include="Microsoft.Rest.ClientRuntime.Azure.TestFramework">
       <HintPath>..\..\..\packages\Microsoft.Rest.ClientRuntime.Azure.TestFramework.1.5.0-preview\lib\net45\Microsoft.Rest.ClientRuntime.Azure.TestFramework.dll</HintPath>
       <Private>True</Private>
@@ -186,6 +186,7 @@
     <Compile Include="Certificates\StopBatchCertificateDeletionCommandTests.cs" />
     <Compile Include="ComputeNodes\DisableBatchComputeNodeSchedulingCommandTests.cs" />
     <Compile Include="ComputeNodes\EnableBatchComputeNodeSchedulingCommandTests.cs" />
+    <Compile Include="ComputeNodes\StartBatchComputeNodeServiceLogUploadCommandTests.cs" />
     <Compile Include="ComputeNodes\GetBatchComputeNodeLoginSettingsCommandTests.cs" />
     <Compile Include="ComputeNodes\RemoveBatchComputeNodeCommandTests.cs" />
     <Compile Include="ComputeNodes\ResetBatchComputeNodeCommandTests.cs" />
@@ -212,6 +213,7 @@
     <Compile Include="Models\BatchAccountContextTest.cs" />
     <Compile Include="Pools\DisableBatchAutoScaleCommandTests.cs" />
     <Compile Include="Pools\EnableBatchAutoScaleCommandTests.cs" />
+    <Compile Include="Pools\GetBatchPoolNodeCountsCommandTests.cs" />
     <Compile Include="Pools\GetBatchPoolCommandTests.cs" />
     <Compile Include="Pools\GetBatchPoolStatisticsCommandTests.cs" />
     <Compile Include="Pools\GetBatchPoolUsageMetricsCommandTests.cs" />

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

@@ -0,0 +1,151 @@
+// ----------------------------------------------------------------------------------
+//
+// 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.Management.Automation;
+using Xunit;
+using BatchClient = Microsoft.Azure.Commands.Batch.Models.BatchClient;
+using ProxyModels = Microsoft.Azure.Batch.Protocol.Models;
+
+namespace Microsoft.Azure.Commands.Batch.Test.ComputeNodes
+{
+    public class StartBatchComputeNodeServiceLogUploadCommandTests : WindowsAzure.Commands.Test.Utilities.Common.RMTestBase
+    {
+        private StartBatchComputeNodeServiceLogUploadCommand startComputeNodeServiceLogUploadCommand;
+        private GetBatchComputeNodeCommand getComputeNodeCommand;
+        private Mock<BatchClient> batchClientMock;
+        private Mock<ICommandRuntime> commandRuntimeMock;
+        private const string fakeUrl = "https://containerUrl?sv==";
+
+        public StartBatchComputeNodeServiceLogUploadCommandTests(Xunit.Abstractions.ITestOutputHelper output)
+        {
+            ServiceManagemenet.Common.Models.XunitTracingInterceptor.AddToContext(new ServiceManagemenet.Common.Models.XunitTracingInterceptor(output));
+            batchClientMock = new Mock<BatchClient>();
+            commandRuntimeMock = new Mock<ICommandRuntime>();
+
+            getComputeNodeCommand = new GetBatchComputeNodeCommand()
+            {
+                CommandRuntime = commandRuntimeMock.Object,
+                BatchClient = batchClientMock.Object,
+            };
+
+            startComputeNodeServiceLogUploadCommand = new StartBatchComputeNodeServiceLogUploadCommand()
+            {
+                CommandRuntime = commandRuntimeMock.Object,
+                BatchClient = batchClientMock.Object,
+            };
+        }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void WhenStartBatchComputeNodeServiceLogUploadCommandIsCalledWithPoolIdAndComputeNodeId_ShouldSucceed()
+        {
+            // Setup cmdlet to get a compute node by id
+            BatchAccountContext context = BatchTestHelpers.CreateBatchContextWithKeys();
+            startComputeNodeServiceLogUploadCommand.BatchContext = context;
+            startComputeNodeServiceLogUploadCommand.PoolId = "pool";
+            startComputeNodeServiceLogUploadCommand.ComputeNodeId = "tvm";
+            startComputeNodeServiceLogUploadCommand.ContainerUrl = fakeUrl;
+            startComputeNodeServiceLogUploadCommand.StartTime = DateTime.UtcNow;
+
+            const int numberOfFilesUploaded = 2;
+            const string virtualDirectoryName = "pool1/tvm";
+
+            // Build a compute node instead of querying the service on a Get ComputeNode call
+            AzureOperationResponse<ProxyModels.UploadBatchServiceLogsResult, ProxyModels.ComputeNodeUploadBatchServiceLogsHeaders> response = 
+                BatchTestHelpers.CreateComputeNodeServiceLogsAddResponse(numberOfFilesUploaded, virtualDirectoryName);
+
+            RequestInterceptor interceptor = BatchTestHelpers.CreateFakeServiceResponseInterceptor<
+                ProxyModels.ComputeNodeUploadBatchServiceLogsOptions,
+                AzureOperationResponse<ProxyModels.UploadBatchServiceLogsResult, ProxyModels.ComputeNodeUploadBatchServiceLogsHeaders>>(response);
+
+            startComputeNodeServiceLogUploadCommand.AdditionalBehaviors = new List<BatchClientBehavior>() { interceptor };
+
+            // Setup the cmdlet to write pipeline output to a list that can be examined later
+            PSStartComputeNodeServiceLogUploadResult result = null;
+            commandRuntimeMock.Setup(r => r.WriteObject(It.IsAny<PSStartComputeNodeServiceLogUploadResult>())).Callback<object>(c => result = (PSStartComputeNodeServiceLogUploadResult)c);
+
+            startComputeNodeServiceLogUploadCommand.ExecuteCmdlet();
+
+            Assert.NotNull(result);
+            Assert.Equal(result.NumberOfFilesUploaded, numberOfFilesUploaded);
+            Assert.Equal(result.VirtualDirectoryName, virtualDirectoryName);
+        }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void WhenStartBatchComputeNodeServiceLogUploadCommandIsCalledWithComputeNode_ShouldSucceed()
+        {
+            // First get a fake tvm
+            BatchAccountContext context = BatchTestHelpers.CreateBatchContextWithKeys();
+            getComputeNodeCommand.BatchContext = context;
+            getComputeNodeCommand.PoolId = "Pool1";
+            getComputeNodeCommand.Id = null;
+            getComputeNodeCommand.Filter = null;
+
+            AzureOperationResponse<IPage<ProxyModels.ComputeNode>, ProxyModels.ComputeNodeListHeaders> response1 =
+                BatchTestHelpers.CreateComputeNodeListResponse(new[] { "tvm1" });
+
+            RequestInterceptor interceptor1 = BatchTestHelpers.CreateFakeServiceResponseInterceptor<
+                ProxyModels.ComputeNodeListOptions,
+                AzureOperationResponse<IPage<ProxyModels.ComputeNode>, ProxyModels.ComputeNodeListHeaders>>(response1);
+
+            var computeNodes = new List<PSComputeNode>();
+            commandRuntimeMock.Setup(r => r.WriteObject(It.IsAny<PSComputeNode>()))
+                .Callback<object>(c => computeNodes.Add((PSComputeNode) c)); 
+
+            getComputeNodeCommand.AdditionalBehaviors = new List<BatchClientBehavior>() { interceptor1 };
+
+            getComputeNodeCommand.ExecuteCmdlet();
+
+            // test StartBatchComputeNodeServiceLogUploadCommand
+            startComputeNodeServiceLogUploadCommand.BatchContext = context;
+            startComputeNodeServiceLogUploadCommand.ComputeNode = computeNodes[0];
+            startComputeNodeServiceLogUploadCommand.ContainerUrl = fakeUrl;
+
+            var utcNow = DateTime.UtcNow;
+            startComputeNodeServiceLogUploadCommand.StartTime = utcNow.AddDays(-1);
+            startComputeNodeServiceLogUploadCommand.EndTime = utcNow;
+
+            const int numberOfFilesUploaded = 2;
+            const string virtualDirectoryName = "pool1/tvm";
+
+            AzureOperationResponse<ProxyModels.UploadBatchServiceLogsResult, ProxyModels.ComputeNodeUploadBatchServiceLogsHeaders> response2 =
+                BatchTestHelpers.CreateComputeNodeServiceLogsAddResponse(numberOfFilesUploaded, virtualDirectoryName);
+
+            RequestInterceptor interceptor2 = BatchTestHelpers.CreateFakeServiceResponseInterceptor<
+                ProxyModels.ComputeNodeUploadBatchServiceLogsOptions,
+                AzureOperationResponse<ProxyModels.UploadBatchServiceLogsResult, ProxyModels.ComputeNodeUploadBatchServiceLogsHeaders>>(response2);
+
+            startComputeNodeServiceLogUploadCommand.AdditionalBehaviors = new List<BatchClientBehavior>() { interceptor2 };
+
+            // Setup the cmdlet to write pipeline output to a list that can be examined later
+            PSStartComputeNodeServiceLogUploadResult result = null;
+            commandRuntimeMock.Setup(r => r.WriteObject(It.IsAny<PSStartComputeNodeServiceLogUploadResult>())).Callback<object>(c => result = (PSStartComputeNodeServiceLogUploadResult)c);
+
+            startComputeNodeServiceLogUploadCommand.ExecuteCmdlet();
+
+            Assert.NotNull(result);
+            Assert.Equal(result.NumberOfFilesUploaded, numberOfFilesUploaded);
+            Assert.Equal(result.VirtualDirectoryName, virtualDirectoryName);
+        }
+    }
+}