Running code in a different session state

July 31, 2023

In the previous post I wrote a function to simplify creating Pester mocks that return values from a provided array.

The function had a few limitations, one of which was that it couldn’t be part of a module. This was briefly mentioned in a comment.

BeforeAll {
  # inline this function (like this) or dot-source it from a script
  # requires modifications to support being inside a module
  function New-SequentialResultsMock {
    ...

Using modules makes it easier to share, deploy and version control our code, so let’s find a way to fix this limitation.

Reproducing the issue

Before we start making changes to the function we need to prove the limitation and have a way to know when it’s fixed. Let’s write a test!

Create the following files in a directory and invoke the test to reproduce the issue.

function New-SequentialResultsMock {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [string] $CommandName,
        [Parameter(Mandatory)]
        [ValidateNotNullOrEmpty()]
        [object[]] $SequentialResults,
        [switch] $Verifiable,
        [ScriptBlock] $ParameterFilter,
        [string] $ModuleName,
        [string[]] $RemoveParameterType,
        [string[]] $RemoveParameterValidation
    )

    # Create the mock's scriptblock in a new closure (dynamic module)
    # to isolate variables. This avoids conflicts between different mocks.
    $mockCallCounter = 0
    $mockScriptBlock = {
        $currentIndex = $script:mockCallCounter++
        if ($mockCallCounter -gt $SequentialResults.Count) {
            throw 'Unexpected call. All out of results.'
        }

        return $SequentialResults[$currentIndex]
    }.GetNewClosure()

    # Remove extra parameters incompatible with Pester\Mock
    $PSBoundParameters.Remove('SequentialResults')

    # Call Mock using the provided parameters and our isolated scriptblock
    Pester\Mock @PSBoundParameters -MockWith $mockScriptBlock
}

Describe 'Sequential Results Mock Helper' {
    Context 'Full Mock helper in a HelperModule' {
        BeforeAll {
            Get-Module HelperModule | Remove-Module
            Import-Module "$PSScriptRoot/HelperModule.psm1"
        }

        It 'Test mock in global session state' {
            New-SequentialResultsMock -CommandName Invoke-Command -SequentialResults @(
                @{ Number = 4; Minimum = Get-Date -Date '2023-06-06 12:00:00Z' },
                @{ Number = 6; Minimum = Get-Date -Date '2023-06-06 12:01:00Z' }
            )

            (Invoke-Command {}).Number | Should -Be 4
            (Invoke-Command {}).Number | Should -Be 6
            { Invoke-Command {} } | Should -Throw 'Unexpected call. All out of results.'
        }
    }
}

> Invoke-Pester ./NewSequentialResultsMock.tests.ps1

Starting discovery in 1 files.
Discovery found 1 tests in 26ms.
Running tests.
[-] Sequential Results Mock Helper.Full Mock helper in a HelperModule.Test mock in global session state 15ms (15ms|1ms)
 Expected 4, but got $null.
 at (Invoke-Command {}).Number | Should -Be 4, /workspaces/Pester/Samples/SequentialMock/NewSequentialResultsMock.tests.ps1:14
 at <ScriptBlock>, /workspaces/Pester/Samples/SequentialMock/NewSequentialResultsMock.tests.ps1:14
Tests completed in 121ms
Tests Passed: 0, Failed: 1, Skipped: 0 NotRun: 0

The test failed as expected. To understand why, we first need to get familiar with the concept of session states.

Session states

In PowerShell session states provides modules with a private space for internal variables, functions and more. Each session state also keeps its own set of scopes. In addition to module-specific session states you also have a default global session state for everything else.

While this isolation is useful it also adds some complexity, especially for scriptblocks as they’re usually bound to (runs in) the session state they were created in. This is how a exported module function can invoke internal functions in the module that you can’t use directly or even see with Get-Command.

While debugging these types of issues it’s important to know which session state you’re currently in. Below you can see an easy way to check this:

# In a fresh PowerShell session
> $ExecutionContext.SessionState | Format-List Module
Module :                # No module = Global session state

# On a breakpoint inside HelperModule.psm1
> $ExecutionContext.SessionState | Format-List Module
Module : HelperModule   # Running in HelperModule's session state

So what went wrong?

When using the Mock-command Pester publishes the new mock in the current session state. You can change this by providing a module name using -ModuleName SomeModuleName. In the failed test Invoke-Command is used directly from the test code which should execute in the global session state, so we didn’t need to override it this time.

However, Mock itself is not invoked from the test code, but from the New-SequentialResultsMock which is now part the HelperModule. This also means the function is executed in the module’s session state. As a result Mock published the Invoke-Command-mock inside the module’s session state, making it invisible in our test code which instead ended up calling the original command.

Usually I would place New-SequentialResultsMock in a ps1-file and dot-source it inside a BeforeAll-block. That would keep everything in the same session state and we’d avoid the problem. This time I promised a fix, so we’ll adjust (over-engineer?) the helper to also work as part of a module.

Fixing the code

Instead of calling Mock directly from our helper, we should do it from the same session state used to call New-SequentialResultsMock, making the helper transparent to Mock. For this to work we need to:

Capture the session state used by the caller
Create a scriptblock to invoke Mock from that is not bound to any specific session state
Invoke the unbound scriptblock in our desired session state

function New-SequentialResultsMock {
26 collapsed lines
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [string] $CommandName,
        [Parameter(Mandatory)]
        [ValidateNotNullOrEmpty()]
        [object[]] $SequentialResults,
        [switch] $Verifiable,
        [ScriptBlock] $ParameterFilter,
        [string] $ModuleName,
        [string[]] $RemoveParameterType,
        [string[]] $RemoveParameterValidation
    )

    # Create the mock's scriptblock in a new closure (dynamic module)
    # to isolate variables. This avoids conflicts between different mocks.
    $mockCallCounter = 0
    $mockScriptBlock = {
        $currentIndex = $script:mockCallCounter++
        if ($mockCallCounter -gt $SequentialResults.Count) {
            throw 'Unexpected call. All out of results.'
        }

        return $SequentialResults[$currentIndex]
    }.GetNewClosure()

    # Remove extra parameters incompatible with Pester\Mock
    $PSBoundParameters.Remove('SequentialResults')

    # Create the mock using the provided parameters and our isolated scriptblock
    Pester\Mock @PSBoundParameters -MockWith $mockScriptBlock
    # Get the session state used by the caller
    $callerSessionState = $PSCmdlet.SessionState

    # Create a scriptblock for invoking Pester's Mock-command.
    # Using `{}.Ast.GetScriptBlock()` to get a scriptblock that is not bound to a specific session state.
    $unboundInvokerScriptBlock = {
        param($BoundParameters, $MockScriptBlock)
        Pester\Mock @BoundParameters -MockWith $MockScriptBlock
    }.Ast.GetScriptBlock()

    # Invoke the unbound scriptblock in the caller session state
    $ExecutionContext.SessionState.InvokeCommand.InvokeScript($callerSessionState, $unboundInvokerScriptBlock, @($PSBoundParameters, $mockScriptBlock))
}

Why not module invoke?

You may have seen code like & (Get-Module MyModule) { 'Run inside a module' } to execute code inside a specific module’s session state. Unfortunately that excludes the global session state which doesn’t have a module we can reference.

PowerShell team-member @SeeminglyScience had (as usual) a workaround in case you prefer this approach:

$ExecutionContext.SessionState.InvokeCommand.InvokeScript($callerSessionState, $unboundInvokerScriptBlock, @($PSBoundParameters, $mockScriptBlock))
# Link the session state to a temporary module-object so we can use normal module invoke method
$tempModule = [psmoduleinfo]::new($false)
$tempModule.SessionState = $callerSessionState
& $tempModule $unboundInvokerScriptBlock -BoundParameters $PSBoundParameters -MockScriptBlock $mockScriptBlock

Making sure it all works

Now it’s time to run the test again. I’ve included an extra test to make sure mocking with -ModuleName also works as expected.

Describe 'Sequential Results Mock Helper' {
    Context 'Full Mock helper in a HelperModule' {
15 collapsed lines
        BeforeAll {
            Get-Module HelperModule | Remove-Module
            Import-Module "$PSScriptRoot/HelperModule.psm1"
        }

        It 'Test mock in global session state' {
            New-SequentialResultsMock -CommandName Invoke-Command -SequentialResults @(
                @{ Number = 4; Minimum = Get-Date -Date '2023-06-06 12:00:00Z' },
                @{ Number = 6; Minimum = Get-Date -Date '2023-06-06 12:01:00Z' }
            )

            (Invoke-Command {}).Number | Should -Be 4
            (Invoke-Command {}).Number | Should -Be 6
            { Invoke-Command {} } | Should -Throw 'Unexpected call. All out of results.'
        }

        It 'Test mock call from inside a module session state' {
            New-SequentialResultsMock -CommandName Invoke-Command -SequentialResults @(
                @{ Number = 4; Minimum = Get-Date -Date '2023-06-06 12:00:00Z' },
                @{ Number = 6; Minimum = Get-Date -Date '2023-06-06 12:01:00Z' }
            ) -ModuleName Pester # Mock only visible to calls from the Pester module

            # Verify mock isn't available in global state
            Invoke-Command {} | Should -BeNullOrEmpty

            # Mock works as expected when called from the Pester module
            InModuleScope Pester {
                (Invoke-Command {}).Number | Should -Be 4
                (Invoke-Command {}).Number | Should -Be 6
                { Invoke-Command {} } | Should -Throw 'Unexpected call. All out of results.'
            }
        }
    }
}

> Invoke-Pester ./NewSequentialResultsMock.tests.ps1

Starting discovery in 1 files.
Discovery found 2 tests in 20ms.
Running tests.
[+] /workspaces/Pester/Samples/SequentialMock/NewSequentialResultsMock.tests.ps1 154ms (80ms|55ms)
Tests completed in 156ms
Tests Passed: 2, Failed: 0, Skipped: 0 NotRun: 0

Everything works! 🎉

Summary

In this post we’ve improved the mock helper-function, making it more portable and easier to maintain and distribute as part of a module.

While working on this we’ve also had a brief introduction to both session states and what it means to have a bound scriptblock. You can find links spread around the article if you’re interested in reading more about these topics.

That’s it for now. Thank you for reading!