MegaBites-AI/Windows-powershell

1# Copyright (c) Microsoft Corporation.2# Licensed under the MIT License.3 4$powershell = Join-Path -Path $PSHOME -ChildPath "pwsh"5 6function Wait-JobPid {7    param (8        $Job9    )10 11    # This is to prevent hanging in the test.12    # Some test environments (such as raspberry_pi) require more time for background job to run.13    $startTime = [DateTime]::Now14    $TimeoutInMilliseconds = 6000015 16    # This will receive the pid of the Job process and nothing more since that was the only thing written to the pipeline.17    do {18        Start-Sleep -Seconds 119        $pwshId = Receive-Job $Job20 21        if (([DateTime]::Now - $startTime).TotalMilliseconds -gt $timeoutInMilliseconds) {22            throw "Unable to receive PowerShell process id."23        }24    } while (!$pwshId)25 26    $pwshId27}28 29# Executes the Enter/Exit PSHostProcess script that returns the pid of the process that's started.30function Invoke-PSHostProcessScript {31    param (32        [string] $ArgumentString,33        [int] $Id,34        [int] $Retry = 5 # Default retry of 5 times35    )36 37    $commandStr = @'38Enter-PSHostProcess {0} -ErrorAction Stop39$PID40Exit-PSHostProcess41'@ -f $ArgumentString42 43    $result = $false44    foreach ($i in 1..$Retry) {45        # use $i as an incrementally growing pause based on the attempt number46        # so that it's more likely to succeed.47        Start-Sleep -Seconds $i48 49        $result = ($commandStr | & $powershell -noprofile -c -) -eq $Id50        if ($result) {51            break52        }53    }54 55    if ($i -gt 1) {56        Write-Verbose -Verbose "Enter-PSHostProcess script failed $i out of $Retry times."57    }58 59    $result60}61 62Describe "Enter-PSHostProcess tests" -Tag Feature {63    Context "By Process Id" {64 65        BeforeAll {66            $oldColor = $env:NO_COLOR67            $env:NO_COLOR = 168        }69 70        AfterAll {71            $env:NO_COLOR = $oldColor72        }73 74        BeforeEach {75            # Start a normal job where the first thing it does is return $PID. After that, spin forever.76            # We will use this job as the target process for Enter-PSHostProcess77            $pwshJob = Start-Job {78                $PID79                while ($true) {80                    Start-Sleep -Seconds 30 | Out-Null81                }82            }83 84            $pwshId = Wait-JobPid $pwshJob85        }86 87        AfterEach {88            $pwshJob | Stop-Job -PassThru | Remove-Job89        }90 91        It "Can enter, exit, and re-enter another PSHost" {92            Wait-UntilTrue { [bool](Get-PSHostProcessInfo -Id $pwshId) } | Should -BeTrue93 94            # This will enter and exit another process95            Invoke-PSHostProcessScript -ArgumentString "-Id $pwshId" -Id $pwshId |96                Should -BeTrue -Because "The script was able to enter another process and grab the pid of '$pwshId'"97 98            # Re-enter and exit the other process99            Invoke-PSHostProcessScript -ArgumentString "-Id $pwshId" -Id $pwshId |100                Should -BeTrue -Because "The script was able to re-enter another process and grab the pid of '$pwshId'."101        }102 103        It "Can enter, exit, and re-enter another Windows PowerShell PSHost" -Skip:(!$IsWindows) {104            # Start a PowerShell job where the first thing it does is return $PID. After that, spin forever.105            # We will use this job as the target process for Enter-PSHostProcess106            $powershellJob = Start-Job -PSVersion 5.1 {107                $PID108                while ($true) {109                    Start-Sleep -Seconds 30 | Out-Null110                }111            }112 113            $powershellId = Wait-JobPid $powershellJob114 115            try {116                Wait-UntilTrue { [bool](Get-PSHostProcessInfo -Id $powershellId) } | Should -BeTrue117 118                # This will enter and exit another process119                Invoke-PSHostProcessScript -ArgumentString "-Id $powershellId" -Id $powershellId |120                    Should -BeTrue -Because "The script was able to enter another process and grab the pid of '$powershellId'."121 122                # Re-enter and exit the other process123                Invoke-PSHostProcessScript -ArgumentString "-Id $powershellId" -Id $powershellId |124                    Should -BeTrue -Because "The script was able to re-enter another process and grab the pid of '$powershellId'."125 126            } finally {127                $powershellJob | Stop-Job -PassThru | Remove-Job128            }129        }130 131        It "Can enter using NamedPipeConnectionInfo" {132            try {133                Wait-UntilTrue { [bool](Get-PSHostProcessInfo -Id $pwshId) } | Should -BeTrue134 135                $npInfo = [System.Management.Automation.Runspaces.NamedPipeConnectionInfo]::new($pwshId)136                $rs = [runspacefactory]::CreateRunspace($npInfo)137 138                # Try to open the runspace while tracing.139                $splat = @{140                    Name = "RunspaceInit"141                    Expression = {$Input.Open()}142                    PSHost = $true143                    ListenerOption = [System.Diagnostics.TraceOptions]::Callstack144                    FilePath = "$TestDrive/$([System.IO.Path]::GetRandomFileName()).log"145                    InputObject = $rs146                }147                Trace-Command @splat148 149                # If opening the runspace fails, then print out the trace with the callstack150                Wait-UntilTrue { $rs.RunspaceStateInfo.State -eq [System.Management.Automation.Runspaces.RunspaceState]::Opened } |151                    Should -BeTrue -Because (Get-Content $splat.FilePath -Raw)152 153                $ps = [powershell]::Create()154                $ps.Runspace = $rs155                $ps.AddScript('$PID')156 157                [int]$retry = 0158                $result = $null159                $errorMsg = "Exception: "160                while ($retry -lt 5 -and $result -eq $null) {161                    try {162                        $result = $ps.Invoke()163                    }164                    catch [System.Management.Automation.Runspaces.InvalidRunspaceStateException] {165                        $errorMsg += $_.Exception.InnerException.Message + "; "166                        $retry++167                        Start-Sleep -Milliseconds 100168                    }169                }170 171                $result | Should -Be $pwshId -Because $errorMsg172            } finally {173                # Clean up disposables174                if ($rs) {175                    $rs.Dispose()176                }177 178                if ($ps) {179                    $ps.Dispose()180                }181            }182        }183    }184 185    Context "By CustomPipeName" {186 187        BeforeAll {188            $oldColor = $env:NO_COLOR189            $env:NO_COLOR = 1190        }191 192        AfterAll {193            $env:NO_COLOR = $oldColor194        }195 196        It "Can enter, exit, and re-enter using CustomPipeName" {197            $pipeName = [System.IO.Path]::GetRandomFileName()198            $pipePath = Get-PipePath -PipeName $pipeName199 200            # Start a job where the first thing it does is set the custom pipe name, then return $PID.201            # After that, spin forever.202            # We will use this job as the target process for Enter-PSHostProcess203            $pwshJob = Start-Job -ArgumentList $pipeName {204                [System.Management.Automation.Remoting.RemoteSessionNamedPipeServer]::CreateCustomNamedPipeServer($args[0])205                $PID206                while ($true) { Start-Sleep -Seconds 30 | Out-Null }207            }208 209            $pwshId = Wait-JobPid $pwshJob210 211            try {212                Wait-UntilTrue { Test-Path $pipePath } | Should -BeTrue213 214                # This will enter and exit another process215                Invoke-PSHostProcessScript -ArgumentString "-CustomPipeName $pipeName" -Id $pwshId |216                    Should -BeTrue -Because "The script was able to enter another process and grab the pipe of '$pipeName'."217 218                # Re-enter and exit the other process219                Invoke-PSHostProcessScript -ArgumentString "-CustomPipeName $pipeName" -Id $pwshId |220                    Should -BeTrue -Because "The script was able to re-enter another process and grab the pipe of '$pipeName'."221 222            } finally {223                $pwshJob | Stop-Job -PassThru | Remove-Job224            }225        }226 227        It "Should throw if CustomPipeName does not exist" {228            { Enter-PSHostProcess -CustomPipeName badpipename } | Should -Throw -ExpectedMessage "No named pipe was found with CustomPipeName: badpipename."229        }230    }231}232