Add new options to Wait-AzureRmHDInsightJob

AzurePowershell.Test.targets

@@ -101,6 +101,7 @@
       <XUnitTests Include=".\src\ResourceManager\RedisCache\Commands.RedisCache.Test\bin\Debug\Microsoft.Azure.Commands.RedisCache.Test.dll"/>
       <!--<XUnitTests Include=".\src\ResourceManager\ServerManagement\Commands.ServerManagement.Test\bin\Debug\Microsoft.Azure.Commands.ServerManagement.Test.dll"/>
       <XUnitTests Include=".\src\ResourceManager\Storage\Commands.Management.Storage.Test\bin\Debug\Microsoft.Azure.Commands.Management.Storage.Test.dll"/>-->
+	  <XUnitTests Include=".\src\ResourceManager\HDInsight\Commands.HDInsight.Test\bin\Debug\Commands.HDInsight.Test.dll"/>
       <XUnitTests Include=".\src\Common\Commands.Common.Authentication.Test\bin\Debug\Microsoft.Azure.Commands.Common.Authentication.Test.dll"/>
       <XUnitTests Include="@(AsmXUnitTests)"/>
     </ItemGroup>

src/ResourceManager/HDInsight/Commands.HDInsight.Test/Commands.HDInsight.Test.csproj

@@ -56,7 +56,7 @@
     </Reference>
     <Reference Include="Microsoft.Azure.Management.HDInsight.Job">
       <SpecificVersion>False</SpecificVersion>
-      <HintPath>..\..\..\packages\Microsoft.Azure.Management.HDInsight.Job.2.0.0-preview\lib\net40\Microsoft.Azure.Management.HDInsight.Job.dll</HintPath>
+      <HintPath>..\..\..\packages\Microsoft.Azure.Management.HDInsight.Job.2.0.3\lib\net40\Microsoft.Azure.Management.HDInsight.Job.dll</HintPath>
     </Reference>
     <Reference Include="Microsoft.Azure.Management.Storage">
       <HintPath>..\..\..\packages\Microsoft.Azure.Management.Storage.2.4.0-preview\lib\net40\Microsoft.Azure.Management.Storage.dll</HintPath>

src/ResourceManager/HDInsight/Commands.HDInsight.Test/ScenarioTests/HDInsightScenarioTestsBase.cs

@@ -44,7 +44,7 @@ protected void SetupManagementClients()
         {
             var hdinsightManagementClient = GetHdInsightManagementClient();
 
-            helper.SetupManagementClients(hdinsightManagementClient);
+            helper.SetupSomeOfManagementClients(hdinsightManagementClient);
         }
 
         protected HDInsightManagementClient GetHdInsightManagementClient()

src/ResourceManager/HDInsight/Commands.HDInsight.Test/UnitTests/JobTests.cs

@@ -12,15 +12,18 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
+using System;
+using System.Collections.Generic;
+using System.Globalization;
+using System.Management.Automation;
+using System.Net;
+using System.Linq;
+using Hyak.Common;
 using Microsoft.Azure.Commands.HDInsight.Models;
 using Microsoft.Azure.Management.HDInsight.Job.Models;
 using Microsoft.WindowsAzure.Commands.Common;
 using Microsoft.WindowsAzure.Commands.ScenarioTest;
 using Moq;
-using System.Collections.Generic;
-using System.Linq;
-using System.Management.Automation;
-using System.Net;
 using Xunit;
 
 namespace Microsoft.Azure.Commands.HDInsight.Test
@@ -379,5 +382,114 @@ public GetAzureHDInsightJobCommand GetJobCommandDefinition()
 
             return cmdlet;
         }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void WaitForJobCompletion_Sucess()
+        {
+            var cmdlet = GetWaitAzureHDInsightJobCommandDefinition(JobCompletionType.Success);
+            cmdlet.ExecuteCmdlet();
+        }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void WaitForJobCompletion_Fail()
+        {
+            var cmdlet = GetWaitAzureHDInsightJobCommandDefinition(JobCompletionType.Fail);
+
+            try
+            {
+                cmdlet.ExecuteCmdlet();
+                throw new Exception("Operation didn't fail");
+            }
+            catch (Exception)
+            {
+            }
+        }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void WaitForJobCompletion_TimeOut()
+        {
+            var expectedExceptionMessage = "HDInsight job with ID {0} has not completed in {1} seconds. Increase the value of -TimeoutInSeconds or check the job runtime.";
+            WaitForJobCompletion_Internal(JobCompletionType.TimeOut, expectedExceptionMessage);
+        }
+
+        public void WaitForJobCompletion_Internal(JobCompletionType completionType, string expectedMessageFormat)
+        {
+            var cmdlet = GetWaitAzureHDInsightJobCommandDefinition(completionType);
+
+            bool success = false;
+
+            try
+            {
+                cmdlet.ExecuteCmdlet();
+            }
+            catch (CloudException ex)
+            {
+                var expectedExceptionMessage = string.Format(CultureInfo.InvariantCulture, expectedMessageFormat, cmdlet.JobId, cmdlet.TimeoutInSeconds);
+                if (ex.Message.Equals(expectedExceptionMessage))
+                {
+                    success = true;
+                }
+            }
+
+            if (!success)
+            {
+                throw new Exception("Operation didn't fail");
+            }
+        }
+
+        public WaitAzureHDInsightJobCommand GetWaitAzureHDInsightJobCommandDefinition(JobCompletionType jobCompetionType)
+        {
+            // Update HDInsight Management properties for Job.
+            SetupManagementClientForJobTests();
+
+            // Setup Job Management mocks
+            const string jobId = "jobid_1984120_001";
+
+            var cmdlet = new WaitAzureHDInsightJobCommand
+            {
+                CommandRuntime = commandRuntimeMock.Object,
+                HDInsightJobClient = hdinsightJobManagementMock.Object,
+                HDInsightManagementClient = hdinsightManagementMock.Object,
+                HttpCredential = new PSCredential("httpuser", string.Format("Password1!").ConvertToSecureString()),
+                ClusterName = ClusterName,
+                JobId = jobId,
+                TimeoutInSeconds = 10
+            };
+
+            JobGetResponse jobDetails = null;
+            TimeSpan? jobDuration = TimeSpan.FromSeconds(10);
+            hdinsightJobManagementMock.Setup(c => c.WaitForJobCompletion(It.IsAny<string>(), It.IsAny<TimeSpan?>(), null))
+                .Callback((string id, TimeSpan? duration, TimeSpan? interval) => { jobDetails = MockJobCompletion(id, jobCompetionType); })
+                .Returns(() => jobDetails);
+
+            return cmdlet;
+        }
+
+        public enum JobCompletionType { Success, Fail, TimeOut };
+
+        public JobGetResponse MockJobCompletion(string jobId, JobCompletionType type)
+        {
+            var jobResponse = new JobGetResponse
+            {
+                JobDetail = new JobDetailRootJsonObject { Id = jobId, Status = new Status(), Userargs = new Userargs() }
+            };
+
+            switch (type)
+            {
+                case JobCompletionType.Success:
+                    break;
+                case JobCompletionType.Fail:
+                    throw new CloudException("Some Failure");
+                case JobCompletionType.TimeOut:
+                    throw new TimeoutException("Some Failure");
+                default:
+                    break;
+            }
+
+            return jobResponse;
+        }
     }
 }

src/ResourceManager/HDInsight/Commands.HDInsight.Test/packages.config

@@ -7,7 +7,7 @@
   <package id="Microsoft.Azure.KeyVault.Core" version="1.0.0" targetFramework="net45" />
   <package id="Microsoft.Azure.Management.Authorization" version="0.18.2-preview" targetFramework="net45" />
   <package id="Microsoft.Azure.Management.HDInsight" version="1.0.14-preview" targetFramework="net45" />
-  <package id="Microsoft.Azure.Management.HDInsight.Job" version="2.0.0-preview" targetFramework="net45" />
+  <package id="Microsoft.Azure.Management.HDInsight.Job" version="2.0.3" targetFramework="net45" />
   <package id="Microsoft.Azure.Management.Resources" version="2.20.0-preview" targetFramework="net45" />
   <package id="Microsoft.Azure.Management.Storage" version="2.4.0-preview" targetFramework="net45" />
   <package id="Microsoft.Azure.Test.Framework" version="1.0.5945.28173-prerelease" targetFramework="net45" />
@@ -34,4 +34,4 @@
   <package id="xunit.extensibility.core" version="2.1.0" targetFramework="net45" />
   <package id="xunit.extensibility.execution" version="2.1.0" targetFramework="net45" />
   <package id="xunit.runner.visualstudio" version="2.1.0" targetFramework="net45" />
-</packages>
+</packages>

src/ResourceManager/HDInsight/Commands.HDInsight/Commands.HDInsight.csproj

@@ -128,7 +128,7 @@
     </Reference>
     <Reference Include="Microsoft.Azure.Management.HDInsight.Job">
       <SpecificVersion>False</SpecificVersion>
-      <HintPath>..\..\..\packages\Microsoft.Azure.Management.HDInsight.Job.2.0.0-preview\lib\net40\Microsoft.Azure.Management.HDInsight.Job.dll</HintPath>
+      <HintPath>..\..\..\packages\Microsoft.Azure.Management.HDInsight.Job.2.0.3\lib\net40\Microsoft.Azure.Management.HDInsight.Job.dll</HintPath>
     </Reference>
     <Reference Include="Microsoft.Azure.ResourceManager, Version=2.0.0.0, Culture=neutral, PublicKeyToken=31bf3856ad364e35, processorArchitecture=MSIL">
       <SpecificVersion>False</SpecificVersion>

src/ResourceManager/HDInsight/Commands.HDInsight/JobCommands/WaitAzureHDInsightJobCommand.cs

@@ -12,11 +12,14 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
+using System;
+using System.Globalization;
+using System.Management.Automation;
 using Hyak.Common;
 using Microsoft.Azure.Commands.HDInsight.Commands;
 using Microsoft.Azure.Commands.HDInsight.Models;
+using Microsoft.Azure.Management.HDInsight.Job.Models;
 using Microsoft.WindowsAzure.Commands.Common;
-using System.Management.Automation;
 
 namespace Microsoft.Azure.Commands.HDInsight
 {
@@ -65,6 +68,12 @@ public PSCredential HttpCredential
         [Parameter(HelpMessage = "The name of the resource group.")]
         public string ResourceGroupName { get; set; }
 
+        [Parameter(HelpMessage = "The total time to wait for job completion, in seconds.")]
+        public int TimeoutInSeconds { get; set; }
+
+        [Parameter(HelpMessage = "The time to wait between job status checks, in seconds.")]
+        public int WaitIntervalInSeconds { get; set; }
+
         #endregion
 
         public override void ExecuteCmdlet()
@@ -75,17 +84,40 @@ public override void ExecuteCmdlet()
             }
 
             var jobDetails = WaitJob();
+
             WriteObject(jobDetails);
         }
 
         public AzureHDInsightJob WaitJob()
         {
             _clusterName = GetClusterConnection(ResourceGroupName, ClusterName);
-            var jobDetail = HDInsightJobClient.GetJob(JobId).JobDetail;
-            while (!jobDetail.Status.JobComplete)
+
+            TimeSpan? duration = null;
+
+            if (TimeoutInSeconds > 0)
+            {
+                duration = TimeSpan.FromSeconds(TimeoutInSeconds);
+            }
+
+            TimeSpan? waitInterval = null;
+
+            if (WaitIntervalInSeconds > 0)
             {
-                jobDetail = HDInsightJobClient.GetJob(JobId).JobDetail;
+                waitInterval = TimeSpan.FromSeconds(WaitIntervalInSeconds);
             }
+
+            JobDetailRootJsonObject jobDetail = null;
+
+            try
+            {
+                jobDetail = HDInsightJobClient.WaitForJobCompletion(JobId, duration, waitInterval).JobDetail;
+            }
+            catch (TimeoutException)
+            {
+                var exceptionMessage = string.Format(CultureInfo.InvariantCulture, "HDInsight job with ID {0} has not completed in {1} seconds. Increase the value of -TimeoutInSeconds or check the job runtime.", JobId, TimeoutInSeconds);
+                throw new CloudException(exceptionMessage);
+            }
+
             var jobDetails = new AzureHDInsightJob(jobDetail, HDInsightJobClient.ClusterName);
             return jobDetails;
         }