refactored common code.

microsoft · Nov 4, 2024 · f41ec54 · f41ec54
1 parent 80834ad
commit f41ec54
Show file tree

Hide file tree

Showing 4 changed files with 63 additions and 95 deletions.
diff --git a/scripts/cleanup_ebpf_cicd_tests.ps1 b/scripts/cleanup_ebpf_cicd_tests.ps1
@@ -14,7 +14,7 @@ Push-Location $WorkingDirectory
 $TestVMCredential = Get-StoredCredential -Target $Target -ErrorAction Stop
 
 # Read the test execution json.
-$TestExecutionConfig = Get-Content ("{0}\{1}" -f $PSScriptRoot, $TestExecutionJsonFileName) | ConvertFrom-Json
+$Config = Get-Content ("{0}\{1}" -f $PSScriptRoot, $TestExecutionJsonFileName) | ConvertFrom-Json
 
 $Job = Start-Job -ScriptBlock {
     param ([Parameter(Mandatory = $True)] [PSCredential] $TestVMCredential,
@@ -64,43 +64,20 @@ $Job = Start-Job -ScriptBlock {
 
 }  -ArgumentList (
     $TestVMCredential,
-    $TestExecutionConfig,
+    $Config,
     $SelfHostedRunnerName,
     $LogFileName,
     $WorkingDirectory,
     $KmTracing)
 
-# Keep track of the last received output count
-$TimeElapsed = 0
-$JobTimedOut = $false
-
-# Loop to fetch and print job output in near real-time
-while ($Job.State -eq 'Running') {
-    $JobOutput = Receive-Job -Job $job
-	$JobOutput | ForEach-Object { Write-Host $_ }
-
-    Start-Sleep -Seconds 2
-    $TimeElapsed += 2
-
-    if ($TimeElapsed -gt $TestJobTimeout) {
-        if ($Job.State -eq "Running") {
-            $VMList = $TestExecutionConfig.VMMap.$SelfHostedRunnerName
-            # currently one VM is used per runner.
-            $TestVMName = $VMList[0].Name
-            Write-Host "Cleaning up VM $TestVMName for Kernel Tests has timed out after one hour" -ForegroundColor Yellow
-            $Timestamp = Get-Date -Format 'yyyy-MM-dd_HH-mm-ss'
-            $CheckpointName = "Cleanup-Hang-$TestVMName-Checkpoint-$Timestamp"
-            Write-Log "Taking snapshot $CheckpointName of $TestVMName"
-            Checkpoint-VM -Name $TestVMName -SnapshotName $CheckpointName
-            $JobTimedOut = $true
-            break
-        }
-    }
-}
 
-# Print any remaining output after the job completes
-$JobOutput = Receive-Job -Job $job
-$JobOutput | ForEach-Object { Write-Host $_ }
+# Wait for the job to complete
+$JobTimedOut = `
+    Wait-TestJobToComplete -Job $Job `
+    -Config $Config `
+    -SelfHostedRunnerName $SelfHostedRunnerName `
+    -TestJobTimeout $TestJobTimeout `
+    -CheckpointPrefix "Cleanup"
 
 # Clean up
 Remove-Job -Job $Job -Force

diff --git a/scripts/common.psm1 b/scripts/common.psm1
@@ -72,4 +72,43 @@ function Compress-File
             break;
         }
     }
+}
+
+function Wait-TestJobToComplete
+{
+    param([Parameter(Mandatory = $true)] [System.Management.Automation.Job] $Job,
+           [Parameter(Mandatory = $true)] [PSCustomObject] $Config,
+           [Parameter(Mandatory = $true)] [string] $SelfHostedRunnerName,
+           [Parameter(Mandatory = $true)] [int] $TestJobTimeout,
+           [Parameter(Mandatory = $true)] [string] $CheckpointPrefix)
+    $TimeElapsed = 0
+    # Loop to fetch and print job output in near real-time
+    while ($Job.State -eq 'Running') {
+        $JobOutput = Receive-Job -Job $job
+        $JobOutput | ForEach-Object { Write-Host $_ }
+
+        Start-Sleep -Seconds 2
+        $TimeElapsed += 2
+
+        if ($TimeElapsed -gt $TestJobTimeout) {
+            if ($Job.State -eq "Running") {
+                $VMList = $Config.VMMap.$SelfHostedRunnerName
+                # currently one VM runs per runner.
+                $TestVMName = $VMList[0].Name            
+                Write-Host "Running kernel tests on $TestVMName has timed out after one hour" -ForegroundColor Yellow
+                $Timestamp = Get-Date -Format 'yyyy-MM-dd_HH-mm-ss'
+                $CheckpointName = "$CheckpointPrefix-$TestVMName-Checkpoint-$Timestamp"
+                Write-Log "Taking snapshot $CheckpointName of $TestVMName"
+                Checkpoint-VM -Name $TestVMName -SnapshotName $CheckpointName
+                $JobTimedOut = $true
+                break
+            }
+        }
+    }
+
+    # Print any remaining output after the job completes
+    $JobOutput = Receive-Job -Job $job
+    $JobOutput | ForEach-Object { Write-Host $_ }
+
+    return $JobTimedOut
 }
diff --git a/scripts/execute_ebpf_cicd_tests.ps1 b/scripts/execute_ebpf_cicd_tests.ps1
@@ -80,36 +80,12 @@ $Job = Start-Job -ScriptBlock {
     $UserModeDumpFolder)
 
 # Keep track of the last received output count
-$TimeElapsed = 0
-$JobTimedOut = $false
-
-# Loop to fetch and print job output in near real-time
-while ($Job.State -eq 'Running') {
-    $JobOutput = Receive-Job -Job $job
-	$JobOutput | ForEach-Object { Write-Host $_ }
-
-    Start-Sleep -Seconds 2
-    $TimeElapsed += 2
-
-    if ($TimeElapsed -gt $TestJobTimeout) {
-        if ($Job.State -eq "Running") {
-            $VMList = $Config.VMMap.$SelfHostedRunnerName
-            # currently one VM runs per runner.
-            $TestVMName = $VMList[0].Name            
-            Write-Host "Running kernel tests on $TestVMName has timed out after one hour" -ForegroundColor Yellow
-            $Timestamp = Get-Date -Format 'yyyy-MM-dd_HH-mm-ss'
-            $CheckpointName = "Execution-Hang-$TestVMName-Checkpoint-$Timestamp"
-            Write-Log "Taking snapshot $CheckpointName of $TestVMName"
-            Checkpoint-VM -Name $TestVMName -SnapshotName $CheckpointName
-            $JobTimedOut = $true
-            break
-        }
-    }
-}
-
-# Print any remaining output after the job completes
-$JobOutput = Receive-Job -Job $job
-$JobOutput | ForEach-Object { Write-Host $_ }
+$JobTimedOut = `
+    Wait-TestJobToComplete -Job $Job `
+    -Config $Config `
+    -SelfHostedRunnerName $SelfHostedRunnerName `
+    -TestJobTimeout $TestJobTimeout `
+    -CheckpointPrefix "Execute"
 
 # Clean up
 Remove-Job -Job $Job -Force

diff --git a/scripts/setup_ebpf_cicd_tests.ps1 b/scripts/setup_ebpf_cicd_tests.ps1
@@ -22,7 +22,7 @@ Import-Module .\common.psm1 -Force -ArgumentList ($LogFileName) -WarningAction S
 Import-Module .\config_test_vm.psm1 -Force -ArgumentList ($TestVMCredential.UserName, $TestVMCredential.Password, $WorkingDirectory, $LogFileName) -WarningAction SilentlyContinue
 
 # Read the test execution json.
-$TestExecutionConfig = Get-Content ("{0}\{1}" -f $PSScriptRoot, $TestExecutionJsonFileName) | ConvertFrom-Json
+$Config = Get-Content ("{0}\{1}" -f $PSScriptRoot, $TestExecutionJsonFileName) | ConvertFrom-Json
 $VMList = $Config.VMMap.$SelfHostedRunnerName
 
 # Delete old log files if any.
@@ -84,45 +84,21 @@ $Job = Start-Job -ScriptBlock {
     Pop-Location
 }  -ArgumentList (
     $TestVMCredential,
-    $TestExecutionConfig,
+    $Config,
     $SelfHostedRunnerName,
     $TestMode,
     $LogFileName,
     $WorkingDirectory,
     $KmTracing,
     $KmTraceType)
 
-# Keep track of the last received output count
-$TimeElapsed = 0
-$JobTimedOut = $false
-
-# Loop to fetch and print job output in near real-time
-while ($Job.State -eq 'Running') {
-    $JobOutput = Receive-Job -Job $job
-	$JobOutput | ForEach-Object { Write-Host $_ }
-
-    Start-Sleep -Seconds 2
-    $TimeElapsed += 2
-
-    if ($TimeElapsed -gt $TestJobTimeout) {
-        if ($Job.State -eq "Running") {
-            $VMList = $Config.VMMap.$SelfHostedRunnerName
-            # currently one VM is used per runner.
-            $TestVMName = $VMList[0].Name
-            Write-Host "Setting up VM $TestVMName for Kernel Tests has timed out after one hour" -ForegroundColor Yellow
-            $Timestamp = Get-Date -Format 'yyyy-MM-dd_HH-mm-ss'
-            $CheckpointName = "Setup-Hang-$TestVMName-Checkpoint-$Timestamp"
-            Write-Log "Taking snapshot $CheckpointName of $TestVMName"
-            Checkpoint-VM -Name $TestVMName -SnapshotName $CheckpointName
-            $JobTimedOut = $true
-            break
-        }
-    }
-}
-
-# Print any remaining output after the job completes
-$JobOutput = Receive-Job -Job $job
-$JobOutput | ForEach-Object { Write-Host $_ }
+# wait for the job to complete
+$JobTimedOut = `
+    Wait-TestJobToComplete -Job $Job `
+    -Config $Config `
+    -SelfHostedRunnerName $SelfHostedRunnerName `
+    -TestJobTimeout $TestJobTimeout `
+    -CheckpointPrefix "Setup"
 
 # Clean up
 Remove-Job -Job $Job -Force