Automatically clean up changes in env variables after test

restore-environment.ps1

# this snippet describes how to automatically clean up
# changes in environment varables
# it can clean up re-writing, adding and removing
# scroll down to see how to use it in real test

Get-Module Pester | Remove-Module 
Import-Module Pester -MaximumVersion 4.9.9
 
function Get-Environment {
    $vars = Get-Item Env:\ 
    $hash = @{}
    foreach ($var in $vars) { 
        $hash.Add($var.Name, $var.Value)
    }
    $hash
}

function Set-Environment { 
    param (
        [Parameter(Mandatory=$true)]
        [Hashtable] $Environment
    )

    # set all items to their previous state
    foreach ($env in $Environment.GetEnumerator()) {
        Set-Item "Env:\$($env.Key)" -Value $env.Value
    }

    $vars = Get-Item "Env:\"
    # remove all items that were not present before
    foreach ($var in $vars) { 
        if (-not $Environment.ContainsKey($var.Name)) {
            Remove-Item "Env:\$($var.Name)"
        }
    }
}

Describe "Restore environment" { 
    It "Restores the values to the state before the test" {
        try {
            # -- Arrange
            $computername = $env:COMPUTERNAME
            $environment = Get-Environment
            $env:COMPUTERNAME = "abc"

            # -- Act
            Set-Environment -Environment $environment

            # -- Assert
            $env:COMPUTERNAME | Should -Be $computername
        }
        finally {
            # in case the code did not fix it we fix it ourselves
            # to avoid changing the computername
            $env:COMPUTERNAME = $computername
        }
    }

    It "Removes all variables that did not exist before the test" {
        try {
            # -- Arrange
            "Env:\abc" | Should -Not -Exist -Because "we are testing that adding the variable will remove it at the end"
            $environment = Get-Environment
            
            $env:abc = "abc"

            # -- Act
            Set-Environment -Environment $environment

            # -- Assert
            "Env:\abc" | Should -Not -Exist
        }
        finally {
            # in case the code did not fix it we fix it ourselves
            # to avoid leaking the test variable
            if (Test-Path "env:\abc") {
                Remove-Item "env:\abc"
            }
        }
    }

    It "Adds all variables that were removed during the test" {
        try {
            # -- Arrange
            $env:abc = "abc"
            "Env:\abc" | Should -Exist -Because "we are testing that removing the variable will add it at the end"
            $environment = Get-Environment
            
            Remove-Item "Env:\abc"

            # -- Act
            Set-Environment -Environment $environment

            # -- Assert
            "Env:\abc" | Should -Exist
        }
        finally {
            # in case the code did not fix it we fix it ourselves
            # to avoid leaking the test variable
            if (Test-Path "env:\abc") {
                Remove-Item "env:\abc"
            }
        }
    }
}

Describe "Testing usage in real test" { 
    # this is how the code would be used in real life 
    Describe "Do test" {
        BeforeEach { 
            $environment = Get-Environment
        }

        It "Adds environment variable" {
            $env:gef = "gef"
        }

        AfterEach {
            Set-Environment $environment
        }
    }

    Describe "Environment variable did not leak" {
        It "env:gef does not exist" {
            "Env:\gef" | Should -Not -Exist
        }

        AfterAll {
            # in case the code did not fix it we fix it ourselves
            # to avoid leaking the test variable
            if (Test-Path "env:\gef") {
                Remove-Item "env:\gef"
            }
        }
    }

}