From c50afe27cd8d3fea12cfe0553f9b1802670136d8 Mon Sep 17 00:00:00 2001
From: Aleksandr Chebotov <47745270+al-cheb@users.noreply.github.com>
Date: Thu, 2 Dec 2021 10:56:43 +0300
Subject: [PATCH] [Windows] Add pause after Windows Update (#4641)

---
 images/win/scripts/ImageHelpers/InstallHelpers.ps1   | 10 +++++-----
 .../Installers/Wait-WindowsUpdatesForInstall.ps1     | 12 ++++++++++++
 images/win/windows2019.json                          |  1 +
 images/win/windows2022.json                          |  4 ++++
 4 files changed, 22 insertions(+), 5 deletions(-)
 create mode 100644 images/win/scripts/Installers/Wait-WindowsUpdatesForInstall.ps1

diff --git a/images/win/scripts/ImageHelpers/InstallHelpers.ps1 b/images/win/scripts/ImageHelpers/InstallHelpers.ps1
index eddac73d..f113071d 100644
--- a/images/win/scripts/ImageHelpers/InstallHelpers.ps1
+++ b/images/win/scripts/ImageHelpers/InstallHelpers.ps1
@@ -519,7 +519,7 @@ function Get-WindowsUpdatesHistory {
             43 {
                 $status = "InProgress"
                 $title = $event.Properties[0].Value
-                break 
+                break
             }
         }
 
@@ -538,7 +538,7 @@ function Invoke-SBWithRetry {
     param (
         [scriptblock] $Command,
         [int] $RetryCount = 10,
-        [int] $TimeoutInSecs = 5
+        [int] $RetryIntervalSeconds = 5
     )
 
     while ($RetryCount -gt 0) {
@@ -547,15 +547,15 @@ function Invoke-SBWithRetry {
             return
         }
         catch {
-            Write-Host "There is an error encounterd:`n $_"
+            Write-Host "There is an error encountered:`n $_"
             $RetryCount--
 
             if ($RetryCount -eq 0) {
                 exit 1
             }
 
-            Write-Host "Waiting $TimeoutInSecs seconds before retrying. Retries left: $RetryCount"
-            Start-Sleep -Seconds $TimeoutInSecs
+            Write-Host "Waiting $RetryIntervalSeconds seconds before retrying. Retries left: $RetryCount"
+            Start-Sleep -Seconds $RetryIntervalSeconds
         }
     }
 }
diff --git a/images/win/scripts/Installers/Wait-WindowsUpdatesForInstall.ps1 b/images/win/scripts/Installers/Wait-WindowsUpdatesForInstall.ps1
new file mode 100644
index 00000000..2039e939
--- /dev/null
+++ b/images/win/scripts/Installers/Wait-WindowsUpdatesForInstall.ps1
@@ -0,0 +1,12 @@
+################################################################################
+##  File:  Wait-WindowsUpdatesForInstall.ps1
+##  Desc:  Wait for installation windows updates to complete
+################################################################################
+
+Invoke-SBWithRetry -RetryCount 10 -RetryIntervalSeconds 120 -Command {
+    $inProgress = Get-WindowsUpdatesHistory | Where-Object Status -eq "InProgress"
+    if ( $inProgress ) {
+        $title = $inProgress.Title -join "`n"
+        throw "InProgress: $title"
+    }
+}
diff --git a/images/win/windows2019.json b/images/win/windows2019.json
index abeec44b..09797cdb 100644
--- a/images/win/windows2019.json
+++ b/images/win/windows2019.json
@@ -274,6 +274,7 @@
             "type": "powershell",
             "pause_before": "2m",
             "scripts": [
+                "{{ template_dir }}/scripts/Installers/Wait-WindowsUpdatesForInstall.ps1",
                 "{{ template_dir }}/scripts/Tests/RunAll-Tests.ps1"
             ]
         },
diff --git a/images/win/windows2022.json b/images/win/windows2022.json
index 948afff7..43175723 100644
--- a/images/win/windows2022.json
+++ b/images/win/windows2022.json
@@ -251,11 +251,15 @@
         },
         {
             "type": "windows-restart",
+            "check_registry": true,
+            "restart_check_command": "powershell -command \"& {if ((-not (Get-Process TiWorker.exe -ErrorAction SilentlyContinue)) -and (-not [System.Environment]::HasShutdownStarted) ) { Write-Output 'Restart complete' }}\"",
             "restart_timeout": "30m"
         },
         {
             "type": "powershell",
+            "pause_before": "2m",
             "scripts": [
+                "{{ template_dir }}/scripts/Installers/Wait-WindowsUpdatesForInstall.ps1",
                 "{{ template_dir }}/scripts/Tests/RunAll-Tests.ps1"
             ]
         },