DynamicParam update

Assignment of session state now comes straight from the mocked function's script block, instead of only when the script block is part of a module.
jamessantiago · Sep 3, 2014 · 3f72ac2 · 3f72ac2
1 parent 73a9860
commit 3f72ac2
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/Functions/Mock.ps1 b/Functions/Mock.ps1
@@ -227,9 +227,11 @@ about_Mocking
                 $dynamicParamStatements = Get-DynamicParamBlock -ScriptBlock $contextInfo.Command.ScriptBlock
                 $dynamicParamScriptBlock = [scriptblock]::Create("$cmdletBinding`r`nparam( $paramBlock )`r`n$dynamicParamStatements")
 
-                if ($contextInfo.Command.ScriptBlock.Module)
+                $sessionStateInternal = Get-ScriptBlockScope -ScriptBlock $contextInfo.Command.ScriptBlock
+
+                if ($null -ne $sessionStateInternal)
                 {
-                    Set-ScriptBlockScope -ScriptBlock $dynamicParamScriptBlock -SessionState $contextInfo.Command.ScriptBlock.Module.SessionState
+                    Set-ScriptBlockScope -ScriptBlock $dynamicParamScriptBlock -SessionStateInternal $sessionStateInternal
                 }
             }
         }