fix: improve error handling for registry file imports and telemetry scheduled tasks

Scripts/Features/Import-RegistryFile.ps1

@@ -11,98 +11,98 @@ function Import-RegistryFile {
         $path
     )
 
-    Write-Host $message
-
-    $usesOfflineHive = $script:Params.ContainsKey("Sysprep") -or $script:Params.ContainsKey("User")
-    $regFilePath = Get-RegistryFilePathForFeature -RegistryKey $path
-
-    if (-not (Test-Path $regFilePath)) {
-        $errorMessage = "Unable to find registry file: $path ($regFilePath)"
-        Write-Host "Error: $errorMessage" -ForegroundColor Red
-        return $false
-    }
-
-    $importScript = {
-        param($targetRegFilePath, $hiveContext)
-
-        if ($script:Params.ContainsKey("WhatIf")) {
-            return (Invoke-RegistryOperationsFromRegFile -RegFilePath $targetRegFilePath)
-        }
-
-        # When the target user's hive is already loaded under their SID, the .reg file's
-        # HKEY_USERS\Default paths won't match. Use the PowerShell registry writer instead,
-        # which remaps Default → SID via Split-RegistryPath.
-        $usePowerShellFallbackOnly = $hiveContext -and [bool]$hiveContext.WasAlreadyLoaded
-
-        if ($usePowerShellFallbackOnly) {
-            $fallbackSucceeded = Invoke-RegistryOperationsFromRegFile -RegFilePath $targetRegFilePath
-            if ($fallbackSucceeded) {
-                Write-Host "The operation completed successfully via PowerShell registry writer."
-            }
-            return $fallbackSucceeded
-        }
-
-        $regResult = Invoke-NonBlocking -ScriptBlock {
-            param($targetRegFilePath)
-            $result = @{
-                Output = @()
-                ExitCode = 0
-                Error = $null
-            }
-
-            try {
-                $global:LASTEXITCODE = 0
-                $output = reg import $targetRegFilePath 2>&1
-                $importExitCode = $LASTEXITCODE
-
-                if ($output) {
-                    $result.Output = @($output)
-                }
-                $result.ExitCode = $importExitCode
-
-                if ($importExitCode -ne 0) {
-                    throw "Registry import failed with exit code $importExitCode for '$targetRegFilePath'"
-                }
-            }
-            catch {
-                $result.Error = $_.Exception.Message
-                $result.ExitCode = if ($LASTEXITCODE -ne 0) { $LASTEXITCODE } else { 1 }
-            }
-
-            return $result
-        } -ArgumentList $targetRegFilePath
-
-        $regOutput = @($regResult.Output)
-        $hasSuccess = ($regResult.ExitCode -eq 0) -and -not $regResult.Error
-
-        if ($regOutput) {
-            foreach ($line in $regOutput) {
-                $lineText = if ($line -is [System.Management.Automation.ErrorRecord]) { $line.Exception.Message } else { $line.ToString() }
-                if ($lineText -and $lineText.Length -gt 0) {
-                    if ($hasSuccess) {
-                        Write-Host $lineText
-                    }
-                    else {
-                        Write-Host $lineText -ForegroundColor Red
-                    }
-                }
-            }
-        }
-
-        if (-not $hasSuccess) {
-            $details = if ($regResult.Error) { $regResult.Error } else { "Exit code: $($regResult.ExitCode)" }
-            Write-Warning "reg import failed for '$path'. Falling back to PowerShell registry writer. Details: $details"
-            $fallbackSucceeded = Invoke-RegistryOperationsFromRegFile -RegFilePath $targetRegFilePath
-            if ($fallbackSucceeded) {
-                Write-Host "The operation completed successfully via PowerShell registry writer."
-            }
-            return $fallbackSucceeded
-        }
-
-        return $true
-    }
-
     try {
+        Write-Host $message
+
+        $usesOfflineHive = $script:Params.ContainsKey("Sysprep") -or $script:Params.ContainsKey("User")
+        $regFilePath = Get-RegistryFilePathForFeature -RegistryKey $path
+
+        if (-not (Test-Path $regFilePath)) {
+            $errorMessage = "Unable to find registry file: $path ($regFilePath)"
+            Write-Host "Error: $errorMessage" -ForegroundColor Red
+            return $false
+        }
+
+        $importScript = {
+            param($targetRegFilePath, $hiveContext)
+
+            if ($script:Params.ContainsKey("WhatIf")) {
+                return (Invoke-RegistryOperationsFromRegFile -RegFilePath $targetRegFilePath)
+            }
+
+            # When the target user's hive is already loaded under their SID, the .reg file's
+            # HKEY_USERS\Default paths won't match. Use the PowerShell registry writer instead,
+            # which remaps Default → SID via Split-RegistryPath.
+            $usePowerShellFallbackOnly = $hiveContext -and [bool]$hiveContext.WasAlreadyLoaded
+
+            if ($usePowerShellFallbackOnly) {
+                $fallbackSucceeded = Invoke-RegistryOperationsFromRegFile -RegFilePath $targetRegFilePath
+                if ($fallbackSucceeded) {
+                    Write-Host "The operation completed successfully via PowerShell registry writer."
+                }
+                return $fallbackSucceeded
+            }
+
+            $regResult = Invoke-NonBlocking -ScriptBlock {
+                param($targetRegFilePath)
+                $result = @{
+                    Output = @()
+                    ExitCode = 0
+                    Error = $null
+                }
+
+                try {
+                    $global:LASTEXITCODE = 0
+                    $output = reg import $targetRegFilePath 2>&1
+                    $importExitCode = $LASTEXITCODE
+
+                    if ($output) {
+                        $result.Output = @($output)
+                    }
+                    $result.ExitCode = $importExitCode
+
+                    if ($importExitCode -ne 0) {
+                        throw "Registry import failed with exit code $importExitCode for '$targetRegFilePath'"
+                    }
+                }
+                catch {
+                    $result.Error = $_.Exception.Message
+                    $result.ExitCode = if ($LASTEXITCODE -ne 0) { $LASTEXITCODE } else { 1 }
+                }
+
+                return $result
+            } -ArgumentList $targetRegFilePath
+
+            $regOutput = @($regResult.Output)
+            $hasSuccess = ($regResult.ExitCode -eq 0) -and -not $regResult.Error
+
+            if ($regOutput) {
+                foreach ($line in $regOutput) {
+                    $lineText = if ($line -is [System.Management.Automation.ErrorRecord]) { $line.Exception.Message } else { $line.ToString() }
+                    if ($lineText -and $lineText.Length -gt 0) {
+                        if ($hasSuccess) {
+                            Write-Host $lineText
+                        }
+                        else {
+                            Write-Host $lineText -ForegroundColor Red
+                        }
+                    }
+                }
+            }
+
+            if (-not $hasSuccess) {
+                $details = if ($regResult.Error) { $regResult.Error } else { "Exit code: $($regResult.ExitCode)" }
+                Write-Warning "reg import failed for '$path'. Falling back to PowerShell registry writer. Details: $details"
+                $fallbackSucceeded = Invoke-RegistryOperationsFromRegFile -RegFilePath $targetRegFilePath
+                if ($fallbackSucceeded) {
+                    Write-Host "The operation completed successfully via PowerShell registry writer."
+                }
+                return $fallbackSucceeded
+            }
+
+            return $true
+        }
+
         if ($usesOfflineHive) {
             # Sysprep targets Default user, User targets the specified user. Logged-in users already have their hive mounted under HKU\<SID>.
             $targetUserName = if ($script:Params.ContainsKey("Sysprep")) { "Default" } else { $script:Params.Item("User") }

Scripts/Features/Telemetry-ScheduledTasks.ps1

@@ -55,7 +55,15 @@ function Disable-TelemetryScheduledTasks {
             $result = Invoke-NonBlocking -ScriptBlock {
             param($path, $name)
             Import-Module ScheduledTasks -ErrorAction SilentlyContinue
-            $taskObj = Get-ScheduledTask -TaskPath $path -TaskName $name -ErrorAction SilentlyContinue
+            try {
+                $taskObj = Get-ScheduledTask -TaskPath $path -TaskName $name -ErrorAction Stop
+            }
+            catch {
+                if ($_.CategoryInfo.Category -eq [System.Management.Automation.ErrorCategory]::ObjectNotFound) {
+                    return @{ Success = $true; Status = 'NotFound' }
+                }
+                return @{ Success = $false; Status = 'Error'; Error = $_.Exception.Message }
+            }
             if (-not $taskObj) {
                 return @{ Success = $true; Status = 'NotFound' }
             }
@@ -121,7 +129,15 @@ function Enable-TelemetryScheduledTasks {
             $result = Invoke-NonBlocking -ScriptBlock {
             param($path, $name)
             Import-Module ScheduledTasks -ErrorAction SilentlyContinue
-            $taskObj = Get-ScheduledTask -TaskPath $path -TaskName $name -ErrorAction SilentlyContinue
+            try {
+                $taskObj = Get-ScheduledTask -TaskPath $path -TaskName $name -ErrorAction Stop
+            }
+            catch {
+                if ($_.CategoryInfo.Category -eq [System.Management.Automation.ErrorCategory]::ObjectNotFound) {
+                    return @{ Success = $true; Status = 'NotFound' }
+                }
+                return @{ Success = $false; Status = 'Error'; Error = $_.Exception.Message }
+            }
             if (-not $taskObj) {
                 return @{ Success = $true; Status = 'NotFound' }
             }

Tests/Import-ConfigToParams.Tests.ps1

@@ -67,6 +67,18 @@ Describe 'Test-ConfigConsistency' {
         Test-ConfigConsistency -Config $config | Should -Match 'Apps entries must be strings'
     }
 
+    It 'reports an error for malformed tweak entries' {
+        $config = [PSCustomObject]@{ Version = '1.0'; Tweaks = @(@{ Value = $true }) }
+
+        Test-ConfigConsistency -Config $config | Should -Match 'Tweaks entries must contain Name and Value properties'
+    }
+
+    It 'reports an error for deployment entries missing a required property' {
+        $config = [PSCustomObject]@{ Version = '1.0'; Deployment = @(@{ Name = 'CreateRestorePoint' }) }
+
+        Test-ConfigConsistency -Config $config | Should -Match 'Deployment entries must contain Name and Value properties'
+    }
+
     It 'reports an error for nonnumeric deployment indexes' {
         $config = [PSCustomObject]@{
             Version = '1.0'

Tests/Import-RegistryFile.Tests.ps1

@@ -26,6 +26,13 @@ Describe 'Import-RegistryFile' {
         Should -Invoke Invoke-NonBlocking -Times 0 -Exactly
     }
 
+    It 'returns false when registry file resolution throws' {
+        Mock Get-RegistryFilePathForFeature { throw 'path resolution failed' }
+
+        Import-RegistryFile -message 'Apply' -path 'feature.reg' | Should -BeFalse
+        Should -Invoke Invoke-NonBlocking -Times 0 -Exactly
+    }
+
     It 'uses the PowerShell writer only in WhatIf mode' {
         $script:Params = @{ WhatIf = $true }
         Import-RegistryFile -message 'Apply' -path 'feature.reg' | Should -BeTrue

Tests/Telemetry-ScheduledTasks.Tests.ps1

@@ -95,6 +95,35 @@ Describe 'Disable-TelemetryScheduledTasks' {
         $result.Error | Should -Match 'access denied'
     }
 
+    It 'returns an error when scheduled-task lookup fails' {
+        Mock Invoke-NonBlocking {
+            param($ScriptBlock, $ArgumentList)
+            $script:taskBlock = $ScriptBlock
+            $script:taskArguments = $ArgumentList
+        }
+        Mock Import-Module {}
+        Mock Get-ScheduledTask { throw 'scheduler unavailable' }
+
+        Disable-TelemetryScheduledTasks
+        $result = & $script:taskBlock @script:taskArguments
+
+        $result.Status | Should -Be 'Error'
+        $result.Error | Should -Match 'scheduler unavailable'
+    }
+
+    It 'treats an absent scheduled task as not found' {
+        Mock Invoke-NonBlocking {
+            param($ScriptBlock, $ArgumentList)
+            $script:taskBlock = $ScriptBlock
+            $script:taskArguments = $ArgumentList
+        }
+        Mock Import-Module {}
+        Mock Get-ScheduledTask { $null }
+
+        Disable-TelemetryScheduledTasks
+        (& $script:taskBlock @script:taskArguments).Status | Should -Be 'NotFound'
+    }
+
     It 'reports <Status> task results' -ForEach @(
         @{ Status = 'Disabled'; Expected = 'Disabled Scheduled Task' }
         @{ Status = 'AlreadyDisabled'; Expected = 'already disabled' }
@@ -175,6 +204,22 @@ Describe 'Enable-TelemetryScheduledTasks' {
         Should -Invoke Enable-ScheduledTask -Times 1 -Exactly -ParameterFilter { $TaskPath -eq '\Microsoft\Windows\Test\' -and $TaskName -eq 'Second' }
     }
 
+    It 'returns an error when scheduled-task lookup fails' {
+        Mock Invoke-NonBlocking {
+            param($ScriptBlock, $ArgumentList)
+            $script:taskBlock = $ScriptBlock
+            $script:taskArguments = $ArgumentList
+        }
+        Mock Import-Module {}
+        Mock Get-ScheduledTask { throw 'scheduler unavailable' }
+
+        Enable-TelemetryScheduledTasks
+        $result = & $script:taskBlock @script:taskArguments
+
+        $result.Status | Should -Be 'Error'
+        $result.Error | Should -Match 'scheduler unavailable'
+    }
+
     It 'reports <Status> task results' -ForEach @(
         @{ Status = 'Enabled'; Expected = 'Enabled Scheduled Task' }
         @{ Status = 'AlreadyEnabled'; Expected = 'already enabled' }