feat: Update Mock-EnvironmentVariable

Improve the performance of setting user environment variables by replacing the [Environment]::SetEnvironmentVariable method with setting the underlying registry directly.
BusHero · Jul 30, 2022 · 51c3326 · 51c3326
1 parent 8456537
commit 51c3326
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 2 deletions.
diff --git a/foo b/foo
diff --git a/src/PesterExtensions/Public/Mock-EnvironmentVariable.ps1 b/src/PesterExtensions/Public/Mock-EnvironmentVariable.ps1
@@ -1,11 +1,35 @@
+function Set-EnvironmentVariable {
+	param (
+		[string]
+		$Variable,
+
+		[string]
+		$Value,
+
+		[EnvironmentVariableTarget]
+		$Target
+	)
+	if ($Target -eq [System.EnvironmentVariableTarget]::Process) {
+		[System.Environment]::SetEnvironmentVariable($Variable, $value, 'Process')
+	}
+	elseif ($Target -eq [System.EnvironmentVariableTarget]::User) {
+		if ($Value) {
+			New-ItemProperty -Path 'HKCU:\Environment' -Name $Variable -Value $Value -Force > $null
+		}
+		else {
+			Remove-ItemProperty -Path 'HKCU:\Environment' -Name $Variable -ErrorAction Ignore
+		}
+	}
+}
+
 function Restore {
 	param (
 		[hashtable]
 		$Backup
 	)
 	foreach ($variable in $Backup.Keys) {
 		foreach ($target in $Backup.$variable.Keys) {
-			[Environment]::SetEnvironmentVariable($variable, $Backup.$variable.$target, $target)
+			Set-EnvironmentVariable -Variable $variable -Value $Backup.$variable.$target -Target $target
 		}
 	}
 }
@@ -29,7 +53,7 @@ function Backup {
 			$OriginalValue = [Environment]::GetEnvironmentVariable($variable, $target)
 			$VariableBackup.Add($target, $OriginalValue)
 			if ($Value) {
-				[Environment]::SetEnvironmentVariable($variable, $Value, $target)
+				Set-EnvironmentVariable -Variable $variable -Value $value -Target $target
 			}
 		}
 		$Backup.Add($variable, $VariableBackup)