Merge pull request #3532 from vexx32/ShouldProcess-cmdlets

(#3318) Add ShouldProcess support to rewritten helper cmdlets
chocolatey · Nov 4, 2024 · 5d0e42d · 5d0e42d
2 parents 2866df9 + c209065
commit 5d0e42d
Show file tree

Hide file tree

Showing 10 changed files with 182 additions and 54 deletions.
diff --git a/src/Chocolatey.PowerShell/Commands/InstallChocolateyPathCommand.cs b/src/Chocolatey.PowerShell/Commands/InstallChocolateyPathCommand.cs
@@ -21,7 +21,7 @@
 
 namespace Chocolatey.PowerShell.Commands
 {
-    [Cmdlet(VerbsLifecycle.Install, "ChocolateyPath")]
+    [Cmdlet(VerbsLifecycle.Install, "ChocolateyPath", SupportsShouldProcess = true, ConfirmImpact = ConfirmImpact.Medium)]
     [OutputType(typeof(void))]
     public class InstallChocolateyPathCommand : ChocolateyCmdlet
     {

diff --git a/src/Chocolatey.PowerShell/Commands/SetEnvironmentVariableCommand.cs b/src/Chocolatey.PowerShell/Commands/SetEnvironmentVariableCommand.cs
@@ -21,7 +21,7 @@
 
 namespace Chocolatey.PowerShell.Commands
 {
-    [Cmdlet(VerbsCommon.Set, "EnvironmentVariable")]
+    [Cmdlet(VerbsCommon.Set, "EnvironmentVariable", SupportsShouldProcess = true, ConfirmImpact = ConfirmImpact.Medium)]
     [OutputType(typeof(void))]
     public sealed class SetEnvironmentVariableCommand : ChocolateyCmdlet
     {

diff --git a/src/Chocolatey.PowerShell/Commands/UninstallChocolateyPathCommand.cs b/src/Chocolatey.PowerShell/Commands/UninstallChocolateyPathCommand.cs
@@ -21,7 +21,7 @@
 
 namespace Chocolatey.PowerShell.Commands
 {
-    [Cmdlet(VerbsLifecycle.Uninstall, "ChocolateyPath")]
+    [Cmdlet(VerbsLifecycle.Uninstall, "ChocolateyPath", SupportsShouldProcess = true, ConfirmImpact = ConfirmImpact.Medium)]
     [OutputType(typeof(void))]
     public class UninstallChocolateyPathCommand : ChocolateyCmdlet
     {

diff --git a/src/Chocolatey.PowerShell/Commands/UpdateSessionEnvironmentCommand.cs b/src/Chocolatey.PowerShell/Commands/UpdateSessionEnvironmentCommand.cs
@@ -20,7 +20,7 @@
 
 namespace Chocolatey.PowerShell.Commands
 {
-    [Cmdlet(VerbsData.Update, "SessionEnvironment")]
+    [Cmdlet(VerbsData.Update, "SessionEnvironment", SupportsShouldProcess = true, ConfirmImpact = ConfirmImpact.Low)]
     [OutputType(typeof(void))]
     public sealed class UpdateSessionEnvironmentCommand : ChocolateyCmdlet
     {

diff --git a/src/Chocolatey.PowerShell/Helpers/EnvironmentHelper.cs b/src/Chocolatey.PowerShell/Helpers/EnvironmentHelper.cs
@@ -139,7 +139,11 @@ public static void SetVariable(PSCmdlet cmdlet, string name, EnvironmentVariable
         {
             if (scope == EnvironmentVariableTarget.Process)
             {
-                Environment.SetEnvironmentVariable(name, value);
+                if (cmdlet.ShouldProcess(name, "Set Process environment variable"))
+                {
+                    Environment.SetEnvironmentVariable(name, value);
+                }
+
                 return;
             }
 
@@ -166,32 +170,38 @@ public static void SetVariable(PSCmdlet cmdlet, string name, EnvironmentVariable
 
                 cmdlet.WriteDebug($"Registry type for {name} is/will be {registryType}");
 
-                if (string.IsNullOrEmpty(value))
-                {
-                    registryKey.DeleteValue(name, throwOnMissingValue: false);
-                }
-                else
+                if (cmdlet.ShouldProcess(name, $"Set {scope} environment variable"))
                 {
-                    registryKey.SetValue(name, value, registryType);
+                    if (string.IsNullOrEmpty(value))
+                    {
+                        registryKey.DeleteValue(name, throwOnMissingValue: false);
+                    }
+                    else
+                    {
+                        registryKey.SetValue(name, value, registryType);
+                    }
                 }
             }
 
             try
             {
-                // Trigger environment refresh in explorer.exe:
-                // 1. Notify all windows of environment block change
-                NativeMethods.SendMessageTimeout(
-                    hWnd: (IntPtr)NativeMethods.HWND_BROADCAST,
-                    Msg: NativeMethods.WM_SETTINGCHANGE,
-                    wParam: UIntPtr.Zero,
-                    lParam: "Environment",
-                    fuFlags: 2,
-                    uTimeout: 5000,
-                    out UIntPtr result);
-
-                // 2. Set a user environment variable making the system refresh
-                var setxPath = string.Format(@"{0}\System32\setx.exe", GetVariable(cmdlet, EnvironmentVariables.SystemRoot, EnvironmentVariableTarget.Process));
-                cmdlet.InvokeCommand.InvokeScript($"& \"{setxPath}\" {EnvironmentVariables.ChocolateyLastPathUpdate} \"{DateTime.Now.ToFileTime()}\"");
+                if (cmdlet.ShouldProcess("Environment variables", "Notify system of changes"))
+                {
+                    // Trigger environment refresh in explorer.exe:
+                    // 1. Notify all windows of environment block change
+                    NativeMethods.SendMessageTimeout(
+                        hWnd: (IntPtr)NativeMethods.HWND_BROADCAST,
+                        Msg: NativeMethods.WM_SETTINGCHANGE,
+                        wParam: UIntPtr.Zero,
+                        lParam: "Environment",
+                        fuFlags: 2,
+                        uTimeout: 5000,
+                        out UIntPtr result);
+
+                    // 2. Set a user environment variable making the system refresh
+                    var setxPath = string.Format(@"{0}\System32\setx.exe", GetVariable(cmdlet, EnvironmentVariables.SystemRoot, EnvironmentVariableTarget.Process));
+                    cmdlet.InvokeCommand.InvokeScript($"& \"{setxPath}\" {EnvironmentVariables.ChocolateyLastPathUpdate} \"{DateTime.Now.ToFileTime()}\"");
+                }
             }
             catch (Exception error)
             {
@@ -221,37 +231,40 @@ public static void UpdateSession(PSCmdlet cmdlet)
                 scopeList.Add(EnvironmentVariableTarget.User);
             }
 
-            foreach (var scope in scopeList)
+            if (cmdlet.ShouldProcess("Current process", "Refresh all environment variables"))
             {
-                foreach (var name in GetVariableNames(scope))
+                foreach (var scope in scopeList)
                 {
-                    var value = GetVariable(cmdlet, name, scope);
-                    if (!string.IsNullOrEmpty(value))
+                    foreach (var name in GetVariableNames(scope))
                     {
-                        SetVariable(cmdlet, name, EnvironmentVariableTarget.Process, value);
+                        var value = GetVariable(cmdlet, name, scope);
+                        if (!string.IsNullOrEmpty(value))
+                        {
+                            SetVariable(cmdlet, name, EnvironmentVariableTarget.Process, value);
+                        }
                     }
                 }
-            }
 
-            // Update PATH, combining both scopes' values.
-            var paths = new string[2];
-            paths[0] = GetVariable(cmdlet, EnvironmentVariables.Path, EnvironmentVariableTarget.Machine);
-            paths[1] = GetVariable(cmdlet, EnvironmentVariables.Path, EnvironmentVariableTarget.User);
+                // Update PATH, combining both scopes' values.
+                var paths = new string[2];
+                paths[0] = GetVariable(cmdlet, EnvironmentVariables.Path, EnvironmentVariableTarget.Machine);
+                paths[1] = GetVariable(cmdlet, EnvironmentVariables.Path, EnvironmentVariableTarget.User);
 
-            SetVariable(cmdlet, EnvironmentVariables.Path, EnvironmentVariableTarget.Process, string.Join(";", paths));
+                SetVariable(cmdlet, EnvironmentVariables.Path, EnvironmentVariableTarget.Process, string.Join(";", paths));
 
-            // Preserve PSModulePath as it's almost always updated by process, preserve it
-            SetVariable(cmdlet, EnvironmentVariables.PSModulePath, EnvironmentVariableTarget.Process, psModulePath);
+                // Preserve PSModulePath as it's almost always updated by process, preserve it
+                SetVariable(cmdlet, EnvironmentVariables.PSModulePath, EnvironmentVariableTarget.Process, psModulePath);
 
-            // Preserve user and architecture
-            if (!string.IsNullOrEmpty(userName))
-            {
-                SetVariable(cmdlet, EnvironmentVariables.Username, EnvironmentVariableTarget.Process, userName);
-            }
+                // Preserve user and architecture
+                if (!string.IsNullOrEmpty(userName))
+                {
+                    SetVariable(cmdlet, EnvironmentVariables.Username, EnvironmentVariableTarget.Process, userName);
+                }
 
-            if (!string.IsNullOrEmpty(architecture))
-            {
-                SetVariable(cmdlet, EnvironmentVariables.ProcessorArchitecture, EnvironmentVariableTarget.Process, architecture);
+                if (!string.IsNullOrEmpty(architecture))
+                {
+                    SetVariable(cmdlet, EnvironmentVariables.ProcessorArchitecture, EnvironmentVariableTarget.Process, architecture);
+                }
             }
         }
     }

diff --git a/tests/helpers/common/Get-WhatIfResult.ps1 b/tests/helpers/common/Get-WhatIfResult.ps1
@@ -0,0 +1,29 @@
+function Get-WhatIfResult {
+    <#
+        .SYNOPSIS
+        Runs a $Command in a new powershell.exe process, and then returns *only*
+        the output lines that are prefixed with 'What if:' which are written as
+        console output.
+    #>
+    [CmdletBinding()]
+    param( 
+        # The script to execute in the new process.
+        [Parameter(Mandatory)]
+        [scriptblock]
+        $Command,
+
+        # Any setup scripts that are required for running. All output from this
+        # script block will be suppressed, if possible.
+        [Parameter()]
+        [scriptblock]
+        $Preamble
+    )
+
+    $commandString = @'
+. {{ {0} }} *>&1 > $null
+& {{ {1} }}
+'@ -f $Preamble, $Command
+
+    powershell -NoProfile -NonInteractive -Command $commandString |
+        Where-Object { $_ -like "What if:*" }
+}
diff --git a/tests/pester-tests/powershell-commands/Install-ChocolateyPath.Tests.ps1 b/tests/pester-tests/powershell-commands/Install-ChocolateyPath.Tests.ps1
@@ -1,4 +1,4 @@
-Describe 'Install-ChocolateyPath helper function tests' -Tags Cmdlets {
+Describe 'Install-ChocolateyPath helper function tests' -Tags InstallChocolateyPath, Cmdlets {
     BeforeAll {
         Initialize-ChocolateyTestInstall
 
@@ -10,7 +10,37 @@
         Remove-Module "chocolateyInstaller" -Force
     }
 
-    Context 'Adding and removing PATH values' -ForEach @(
+    Context 'Unit tests' -Tags WhatIf -ForEach @(
+        @{ Scope = 'Process' }
+        @{ Scope = 'User' }
+        @{ Scope = 'Machine' }
+    ) {
+        Context 'Path "<_>"' -ForEach @("C:\test", "C:\tools") {
+            BeforeAll {    
+                $Preamble = [scriptblock]::Create("Import-Module '$testLocation\helpers\chocolateyInstaller.psm1'")
+            }
+
+            It 'stores the value in the desired PATH scope' {
+                $Command = [scriptblock]::Create("Install-ChocolateyPath -Path '$_' -Scope $Scope -WhatIf")
+
+                $results = @( Get-WhatIfResult -Preamble $Preamble -Command $Command )
+                $results[0] | Should -BeExactly "What if: Performing the operation ""Set $Scope environment variable"" on target ""PATH""."
+
+                if ($Scope -ne 'Process') {
+                    $results[1] | Should -BeExactly 'What if: Performing the operation "Notify system of changes" on target "Environment variables".'
+                    $results[2] | Should -BeExactly 'What if: Performing the operation "Refresh all environment variables" on target "Current process".'
+                }
+            }
+
+            It 'skips adding the value if it is already present' {
+                $targetPathEntry = [Environment]::GetEnvironmentVariable('PATH', $Scope) -split ';' | Select-Object -First 1
+                $Command = [scriptblock]::Create("Install-ChocolateyPath -Path '$targetPathEntry' -Scope $Scope -WhatIf")
+                Get-WhatIfResult -Preamble $Preamble -Command $Command | Should -BeNullOrEmpty -Because 'we should skip adding values that already exist'
+            }
+        }
+    }
+
+    Context 'Adding and removing PATH values' -Tag VMOnly -ForEach @(
         @{ Scope = 'Process' }
         @{ Scope = 'User' }
         @{ Scope = 'Machine' }
@@ -55,7 +85,7 @@
     }
 }
 
-Describe 'Install-ChocolateyPath end-to-end tests with add-path package modifying <Scope> PATH' -Tags Cmdlet -ForEach @(
+Describe 'Install-ChocolateyPath end-to-end tests with add-path package modifying <Scope> PATH' -Tags Cmdlet, UninstallChocolateyPath, VMOnly -ForEach @(
     @{ Scope = 'User' }
     @{ Scope = 'Machine' }
 ) {

diff --git a/tests/pester-tests/powershell-commands/Set-EnvironmentVariable.Tests.ps1 b/tests/pester-tests/powershell-commands/Set-EnvironmentVariable.Tests.ps1
@@ -1,11 +1,31 @@
-Describe 'Set-EnvironmentVariable helper function tests' -Tags Cmdlets {
+Describe 'Set-EnvironmentVariable helper function tests' -Tags SetEnvironmentVariable, Cmdlets {
     BeforeAll {
         Initialize-ChocolateyTestInstall
 
         $testLocation = Get-ChocolateyTestLocation
         Import-Module "$testLocation\helpers\chocolateyInstaller.psm1"
     }
 
+    Context 'Unit tests' -Tags WhatIf -ForEach @(
+        @{ Scope = 'Process' }
+        @{ Scope = 'User' }
+        @{ Scope = 'Machine' }
+    ) {
+        It 'Sets an environment variable value at the target <Scope>' {
+            $testVariableName = 'testVariable'
+            $Preamble = [scriptblock]::Create("Import-Module '$testLocation\helpers\chocolateyInstaller.psm1'")
+            $Command = [scriptblock]::Create("Set-EnvironmentVariable -Name $testVariableName -Value 'TEST' -Scope $Scope -WhatIf")
+
+            $results = @( Get-WhatIfResult -Preamble $Preamble -Command $Command )
+            $results[0] | Should -BeExactly "What if: Performing the operation ""Set $Scope environment variable"" on target ""testVariable""."
+
+            if ($Scope -ne 'Process') {
+                $results[1] | Should -BeExactly 'What if: Performing the operation "Notify system of changes" on target "Environment variables".'
+                $results[2] | Should -BeExactly 'What if: Performing the operation "Refresh all environment variables" on target "Current process".'
+            }
+        }
+    }
+
     Context 'Sets an environment variable value at the target <Scope>' -ForEach @(
         @{ Scope = 'Process' }
         @{ Scope = 'User' }

diff --git a/tests/pester-tests/powershell-commands/Uninstall-ChocolateyPath.Tests.ps1 b/tests/pester-tests/powershell-commands/Uninstall-ChocolateyPath.Tests.ps1
@@ -1,12 +1,38 @@
-Describe 'Uninstall-ChocolateyPath helper function tests' -Tags Cmdlets {
+Describe 'Uninstall-ChocolateyPath helper function tests' -Tags UninstallChocolateyPath, Cmdlets {
     BeforeAll {
         Initialize-ChocolateyTestInstall
 
         $testLocation = Get-ChocolateyTestLocation
         Import-Module "$testLocation\helpers\chocolateyInstaller.psm1"
     }
 
-    Context 'Adding and removing PATH values' -ForEach @(
+    Context 'Unit tests' -Tags WhatIf -ForEach @(
+        @{ Scope = 'Process' }
+        @{ Scope = 'User' }
+        @{ Scope = 'Machine' }
+    ) {
+        It 'removes a stored PATH value in the desired PATH scope' {
+            $targetPathEntry = [Environment]::GetEnvironmentVariable('PATH', $Scope) -split ';' | Select-Object -First 1
+            $Preamble = [scriptblock]::Create("Import-Module '$testLocation\helpers\chocolateyInstaller.psm1'")
+            $Command = [scriptblock]::Create("Uninstall-ChocolateyPath -Path '$targetPathEntry' -Scope $Scope -WhatIf")
+
+            $results = @( Get-WhatIfResult -Preamble $Preamble -Command $Command )
+            $results[0] | Should -BeExactly "What if: Performing the operation ""Set $Scope environment variable"" on target ""PATH""."
+
+            if ($Scope -ne 'Process') {
+                $results[1] | Should -BeExactly 'What if: Performing the operation "Notify system of changes" on target "Environment variables".'
+                $results[2] | Should -BeExactly 'What if: Performing the operation "Refresh all environment variables" on target "Current process".'
+            }
+        }
+
+        It 'skips removing the value if it is not present' {
+            $targetPathEntry = [Environment]::GetEnvironmentVariable('PATH', $Scope) -split ';' | Select-Object -First 1
+            $Command = [scriptblock]::Create("Uninstall-ChocolateyPath -Path 'C:\ThisShouldNotBePresent' -Scope $Scope -WhatIf")
+            Get-WhatIfResult -Preamble $Preamble -Command $Command | Should -BeNullOrEmpty -Because 'we should skip removing a value that does not exist'
+        }
+    }
+
+    Context 'Adding and removing PATH values' -Tags VMOnly -ForEach @(
         @{ Scope = 'Process' }
         @{ Scope = 'User' }
         @{ Scope = 'Machine' }
@@ -40,7 +66,7 @@
     }
 }
 
-Describe 'Uninstall-ChocolateyPath end-to-end tests with add-path package' -Tags Cmdlet -ForEach @(
+Describe 'Uninstall-ChocolateyPath end-to-end tests with add-path package' -Tags Cmdlet, UninstallChocolateyPath, VMOnly -ForEach @(
     @{ Scope = 'User' }
     @{ Scope = 'Machine' }
 ) {

diff --git a/tests/pester-tests/powershell-commands/Update-SessionEnvironment.Tests.ps1 b/tests/pester-tests/powershell-commands/Update-SessionEnvironment.Tests.ps1
@@ -1,10 +1,20 @@
-Describe 'Update-SessionEnvironment helper function tests' -Tag Cmdlets {
+Describe 'Update-SessionEnvironment helper function tests' -Tag UpdateSessionEnvironment, Cmdlets {
     BeforeAll {
         Initialize-ChocolateyTestInstall
 
         $testLocation = Get-ChocolateyTestLocation
         Import-Module "$testLocation\helpers\chocolateyInstaller.psm1"
     }
+
+    Context 'Unit tests' -Tag WhatIf {
+        It 'refreshes the current session environment variables' {
+            $Preamble = [scriptblock]::Create("Import-Module '$testLocation\helpers\chocolateyInstaller.psm1'")
+            $Command = [scriptblock]::Create("Update-SessionEnvironment -WhatIf")
+
+            $results = Get-WhatIfResult -Preamble $Preamble -Command $Command
+            $results | Should -BeExactly 'What if: Performing the operation "refresh all environment variables" on target "current process".'
+        }
+    }
 
     Context 'Refreshing environment' {
         BeforeAll {