Test components and example use for testing PowerShell behavior of assignments, flow control, etc during job stopping

atomicAssignmentSourceMock.cs

using System;
using System.Management.Automation;
using System.Threading;

// Builder type for the arguments to AtomicAssignmentSourceMock.
// The builder pattern is necessary because the values for Job and PipelineStopping only become available in different threads at different time, and only after a reference to an arguments object is already required.
public class AtomicAssignmentSourceMockArgs {
    // flag indicating whether the mock should behave as though this is a dry run
    public bool DryRun = false;

    // Event that is raised when Job is populated
    public ManualResetEventSlim JobSet = new ManualResetEventSlim();
    Job2 _job;

    // The Job property.  Can only be assigned once.  Cannot be assigned nulled.  Raises JobSet event when assigned.
    public Job2 Job {
        get {return _job;}
        set {
            if (null != _job) {
                throw new ArgumentException("Job already set");
            }
            if (null == value) {
                throw new ArgumentNullException("Cannot set Job to null.");
            }
            _job = value;
            JobSet.Set();
        }
    }
    // Event that is raised when PipelineStopping is populated
    public ManualResetEventSlim PipelineStoppingSet = new ManualResetEventSlim();
    ManualResetEventSlim _pipelineStopping;

    // The PipelineStopping property.  Can only be assigned once.  Cannot be assigned nulled. Raises PipelineStoppingSet event when assigned.
    public ManualResetEventSlim PipelineStopping {
        get {return _pipelineStopping;}
        set {
            if (null != _pipelineStopping) {
                throw new ArgumentException("PipelineStopping already set");
            }
            if (null == value) {
                throw new ArgumentNullException("Cannot set PipelineStopping to null.");
            }
            _pipelineStopping = value;
            PipelineStoppingSet.Set();
        }
    }
    public AtomicAssignmentSourceMockArgs() {}
}

// Binary PowerShell Cmdlet for continuing the build of AtomicAssignmentSourceMockArgs.
// This Cmdlet is necessary to synchronize the completion of AtomicAssignmentSourceMock's constructor with the job stopping.  The Cmdlet populates the PipelineProperty of the AtomicAssignmentSourceMockArgs with an event that is raised when the StopProcessing() is invoked on the Cmdlet.  The AtomicAssignmentSourceMockArgs is output using WriteOutput().  Note that the the instance of the Cmdlet must remain in the active pipeline for the duration of the AtomicAssignmentSourceMock constructor call that accepts the AtomicAssignmentSourceMockArgs.
[Cmdlet(VerbsOther.Use,"AtomicAssignmentSourceMockArgs")]
public class UseAtomicAssignmentSourceMockArgsCommand : Cmdlet {
    // The event that is set when StopProcessing() is invoked.
    ManualResetEventSlim _event;

    [Parameter(Mandatory=true)]
    // AtomicAssignmentSourceMockArgs without the PipelineStopping property set.
    public AtomicAssignmentSourceMockArgs Arguments { get; set;}
    protected override void BeginProcessing() {
        // prepare the event the will signal pipeline stopping
        _event = new ManualResetEventSlim();

        // assign the event to the arguments for use downstream in the pipeline
        Arguments.PipelineStopping = _event;

        if (!Arguments.DryRun) {
            // In the case where this is not a dry run, the Job argument is awaited for population here.  The Job argument is required by AtomicAssignmentSourceMock constructor when DryRun is true.
            Arguments.JobSet.Wait();
        }
        // Write the arguments to the pipeline.
        // This accomplishes two things:
        // 1. The arguments are made available at the output of the Cmdlet.
        // 2. The lifetime of a scriptblock downstream of the Cmdlet in the pipeline will keep this instance active.  Such a scriptblock can construct AtomicAssignmentSourceMock with the output arguments.  In that way Arguments.Job.JobStopAsync() and, in turn, StopProcessing() will be called while this command is active.
        WriteObject(Arguments);
    }
    // StopProcessing() is called subsequent to the AtomicAssignmentSourceMock constructor invoking arguments.Job.JobStopAsync().
    protected override void StopProcessing() {
        // Raise the pipeline stopping event.
        _event.Set();
    }
}

// A mock object that synchronizes with the job it is running in to do the following:
// 1. Enter the constructor.
// 2. Stop the job it is running in.
// 3. Wait for a command in a running pipeline to have StopProcessing() invoked.
public class AtomicAssignmentSourceMock {
    public AtomicAssignmentSourceMock (
        bool                 dryRun,
        Job2                 job,
        ManualResetEventSlim pipelineStopping
    ) {
        if (dryRun) {
            // do nothing unusual on a dry run
            return;
        }
        // this is not a dry run, so synchronization must take place

        if (job.JobStateInfo.State != JobState.Running) {
            throw new InvalidOperationException("Job must be Running.");
        }
        // The job must be in the running state here.
        // Stop the job.  This call to StopJobAsync() stops the very job the where scriptblock that calls this constructor is running.
        job.StopJobAsync();
        // Wait for the pipeline to begin stopping subsequent to the StopJobAsync().  Once the pipeline has begun stopping, the constructor can complete, and, if the assignment operation is atomic, the assignment of this new object will complete.  If the assignment operation is not atomic, the assignment of this new object will not complete.
        pipelineStopping.Wait();
    }
    public AtomicAssignmentSourceMock (
        AtomicAssignmentSourceMockArgs arguments) :
        this (arguments.DryRun,arguments.Job,arguments.PipelineStopping) {}
}

atomicAssignmentTest.ps1

#Requires -Module Microsoft.PowerShell.ThreadJob
#Requires -Version 7.3

# add the .Net mock types and Use-AtomicAssignmentSourceMockArgs
Add-Type -Path .\atomicAssignmentSourceMock.cs

@(
    # specify assignment scriptblocks to test
    {$a = [AtomicAssignmentSourceMock]::new($mockArgs)}
    {$a = ([AtomicAssignmentSourceMock]::new($mockArgs))}
    {$a = $([AtomicAssignmentSourceMock]::new($mockArgs))}
    {$a = @([AtomicAssignmentSourceMock]::new($mockArgs))}
    {$a = ($x =[AtomicAssignmentSourceMock]::new($mockArgs))}
    {$a = $(($x =[AtomicAssignmentSourceMock]::new($mockArgs)))}
    {$a = @(($x =[AtomicAssignmentSourceMock]::new($mockArgs)))}
    {$a = [AtomicAssignmentSourceMock]::new($mockArgs),
          [AtomicAssignmentSourceMock]::new($mockArgs)}
    {$a = @([AtomicAssignmentSourceMock]::new($mockArgs)
            [AtomicAssignmentSourceMock]::new($mockArgs))}
    {Set-Variable a ([AtomicAssignmentSourceMock]::new($mockArgs))}
    {$a = New-Object AtomicAssignmentSourceMock $mockArgs}
    {$a = . {[AtomicAssignmentSourceMock]::new($mockArgs)}}
    {$a = & {[AtomicAssignmentSourceMock]::new($mockArgs)}}
) |
ForEach-Object {

    $scriptblock = $_

    # define a scriptblock the completes one run of the test
    $run = {
        param(
            [Parameter(Mandatory)]
            [bool]
            # True  : The mock constructor behaves as normal.
            # False : The mock constructor invokes .StopJobAsync() then waits for .StopProcessing() to be called on Use-AtomicAssignmentSourceMockArgs
            $DryRun,

            [Parameter(Mandatory)]
            [scriptblock]
            # The assignment script block under test.  $a is the variable to which assignment is expected.  The AtomicAssignmentSourceMockArgs are available in $mockArgs.
            #
            # Example:
            #    {$a = [AtomicAssignmentSourceMock]::new($mockArgs)}
            $AssignmentScriptBlock
        )

        # a thread-safe collection for logging
        $log = [System.Collections.Concurrent.ConcurrentQueue[string]]::new()

        # start building the mock arguments
        $mockArgs = [AtomicAssignmentSourceMockArgs]::new();

        # configure whether the mock should treat this as a dry run
        $mockArgs.DryRun = $DryRun

        # Start the job where AssignmentScriptBlock will be executed.
        # This job will be stopped when by [AtomicAssignmentSourceMock]::new() when DryRun is false
        $job =
            Start-ThreadJob                                         `
                -ArgumentList $mockArgs,$log,$AssignmentScriptBlock `
                -ScriptBlock {
                    param($mockArgs,$log,$AssignmentScriptBlock)
                    begin {
                        # import Use-AtomicAssignmentSourceMockArgs for use
                        [UseAtomicAssignmentSourceMockArgsCommand] |
                            % Assembly                       |
                            Import-Module
                    }
                    end {
                        # The building of $mockArgs is continued by this call to Use-AtomicAssignmentSourceMockArgs.  The call here adds the PipelineStopping event to $mockArgs which is raised when .StopProcessing() is invoked on this instance of Use-AtomicSourceMockArgs.  The modified AtomicAssignmentSourceMockArgs are output from Use-AtomicAssignmentSourceMockArgs with WriteObject().  Note that this call to Use-AtomicAssignmentSourceMockArgs must remain a part of the active pipeline for the duration of the assignment under test otherwise .StopProcessing() will not be called and the .PipelineStopping event will not be raised.  The natural way to achieve this is to call Use-AtomicAssignmentSourceMockArgs upstream from the call site of the assignment statement under test.  That is done here.
                        Use-AtomicAssignmentSourceMockArgs `
                                -Arguments $mockArgs |
                            . {
                                process {
                                <#
                                At this stage Use-AtomicAssignmentSourceMockArgs is in the active pipeline, and $mockArgs has its .PipelineStopping property populated.  Considering the value of

                                {$a = [AtomicAssignmentSourceMock]::new($mockArgs)}

                                $AssignmentScriptBlock as an example, the next line will do one of the following:

                                DryRun true  : The constructor will complete immediately without side-effects.
                                DryRun false : The constructor will wait for the .Job property of $mockArgs to be populated.  When .Job is available, the constructor invokes .Job.StopJobAsync() thereby causing the job in which this is being constructed to begin stopping.  The constructor then waits for .PipelineStopping to be raises.  .PipelineStopping is raised when the PowerShell's job stopping machinery eventually calls .StopProcessing() on the instance of Use-AtomicAssignmentSourceMockArgs that is currently upstream in the pipeline.  When .PipelineStopping is raised, the constructor completes.  In this way, the constructor is entered before the job is stopped but completes after the job is stopped.  Whether the assignment completes depends on the particulars of the statement.

                                If the assignment in the example $AssignmentScriptBlock, does not complete, then $a will $null in the clean{} block.  If the assignment does complete, then $a will be not be $null in the clean{}.
                                #>
                                . ($AssignmentScriptBlock.Ast.GetScriptBlock())
                            }}
                        # Record whether the end block completes to flag whether job stopping occurred. If job stopping occurred amidst invoking $AssignmentScriptBlock, then this statement is not reached.
                        $log.Enqueue('end{}')
                    }
                    clean {
                        # Record that the clean{} block was reached.
                        $log.Enqueue('clean{}')

                        # Record whether $a is null.   If $a is null here, it was not assigned to.  This record is used to check for assignment.
                        $log.Enqueue("`$null -eq `$a: $($null -eq $a)")
                    }
            }

        # Continue building the mock arguments.  Note that this might happen either before or after the AtomicAssignmentSourceMock constructor is entered.  That constructor is aware of that possibility and awaits the .Job property to be set when that happens.
        $mockArgs.Job = $job

        # wait for the job to complete, then remove it
        $job | Wait-Job | Remove-Job | Out-Null

        # output the log
        $log
    }

    <# Here the same assignment scriptblock under test is tested twice:
    DryRun true  : The dry run does not involve any job stopping and is used to confirm that assignment succeeds under such normal conditions.
    DryRun false : The real run completes all of the waiting and stopping during the execution of the constructor required to determine whether assignment is atomic.
    #>
    $dryRunLog  = . $run -DryRun $true  -AssignmentScriptBlock $scriptblock
    $realRunLog = . $run -DryRun $false -AssignmentScriptBlock $scriptblock

    # extract whether assignment succeeded in the dry and real runs
    $dry_run_assigns  = [bool]($dryRunLog -eq '$null -eq $a: False')
    $real_run_assigns = [bool]($realRunLog -eq '$null -eq $a: False')

    # prepare the test result for showing on the command line
    [pscustomobject]@{
        # Report whether assignment is atomic.  If assignment failed during dry run, show nothing.
        atomic      = if (-not $dry_run_assigns) {} else {$real_run_assigns}
        # show the literal text of the scriptblock with braces
        scriptblock = "{ $scriptblock }"
    }
} |
    # Show the output as an appropriately-formatted table
    Format-Table `
        -Wrap    `
        -Property @{Name = 'atomic'     ; Expression = 'atomic'     ; Width = 'atomic'.Length},
                  @{Name = 'scriptblock'; Expression = 'scriptblock'; Width = 80             }

exampleTest.ps1

#Requires -Module Microsoft.PowerShell.ThreadJob

# load the mock type
Add-Type -Path .\atomicAssignmentSourceMock.cs

# set up the log
$log = [System.Collections.Concurrent.ConcurrentQueue[string]]::new()

# begin building the mock arguments
$mockArgs = [AtomicAssignmentSourceMockArgs]::new()

# set whether this will be a real or dry run
$mockArgs.DryRun = $false

$job =
    Start-ThreadJob                  `
        -ArgumentList $log,$mockArgs `
        -ScriptBlock {
            param($log,$mockArgs)
            # import Use-AtomicAssignmentSourceMockArgs for use
            [UseAtomicAssignmentSourceMockArgsCommand] |
                % Assembly                             |
                Import-Module

            # Continue to build the mock arguments
            # This command populates the PipelineStopping event so that completion of mock construction can be synchronized with the job stopping.
            Use-AtomicAssignmentSourceMockArgs `
                    -Arguments $mockArgs       |
                . {
                    <#
                    This scriptblock is set up to synchronize the next completion of construction
                    of an AtomicAssignmentSourceMock with job Stopping.

                    Subject this script block to test scenarios involving .StopProcessing() and Stopping to determine the behavior of PowerShell.

                    Use $log.Enqueue() to write to the log that is output to the console after the job is stopped.
                    #>

                    ### Example Test ###
                    ### The following example demonstrates (as of PowerShell 7.4.6) that a clean{} block might not complete
                    begin {}
                    clean {
                        # nothing unusual encountered at this point
                        # log entry into the clean{} block
                        $log.Enqueue('clean{}')

                        # Construct the mock.
                        # $mockArgs.DryRun = $true  : nothing unusual happens during object construction
                        # $mockArgs.DryRun = $false : construction is begun but waits until the Stopping signal reaches the pipeline.
                        $null = [AtomicAssignmentSourceMock]::new($mockArgs)

                        # This line is always executed because, apparently, the Stopping flag is not checked amidst a series of only .Net method calls.
                        $log.Enqueue('clean{} mock constructed')

                        # $mockArgs.DryRun = $false : This next line is not executed because, apparently, Stopping is checked by PowerShell before that scriptblock is invoked
                        $null = . {}

                        # $mockArgs.DryRun = $true  : this line is executed
                        # $mockArgs.DryRun = $false : this line is not executed because execution of this block stopped on the previous line
                        $log.Enqueue('clean{} completed')
                    }
                    ### End Example Test ###
                }
        }

# Continue to build the mock arguments.
$mockArgs.Job = $job

# Wait for the job and output the results.
$job | Receive-Job -Wait

# Output the log.
$log