From ec77ceab7ba99bec93376ad6ca1f6f2d5dfc9d22 Mon Sep 17 00:00:00 2001
From: Matt Wrock <matt@mattwrock.com>
Date: Tue, 16 Oct 2012 07:58:37 -0700
Subject: [PATCH] mock works with any command (not just functions) and the
 parameter filters are implemented. also implemented Clear-Mock

---
 Functions/Mock.Tests.ps1 | 50 ++++++++++++++++++++++++++++++++++++++--
 Functions/Mock.ps1       | 41 +++++++++++++++++++++++++++-----
 2 files changed, 83 insertions(+), 8 deletions(-)

diff --git a/Functions/Mock.Tests.ps1 b/Functions/Mock.Tests.ps1
index 35aca5eba..3f02462ff 100644
--- a/Functions/Mock.Tests.ps1
+++ b/Functions/Mock.Tests.ps1
@@ -2,7 +2,7 @@ $here = Split-Path -Parent $MyInvocation.MyCommand.Path
 $sut = (Split-Path -Leaf $MyInvocation.MyCommand.Path).Replace(".Tests.", ".")
 . "$here\$sut"
 
-function FunctionUnderTest {
+function FunctionUnderTest ([string]$param1=""){
     return "I am a real world test"
 }
 
@@ -19,8 +19,31 @@ Describe "When calling Mock on existing function" {
     It "Should Invoke the mocked script" {
         $result.should.be("I am the mock test")
     }
+    Clear-Mocks
 }
 
+Describe "When calling Mock on existing cmdlet" {
+    Mock Get-Process {return "I am not Get-Process"}
+
+    $result=Get-Process
+
+    It "Should Invoke the mocked script" {
+        $result.should.be("I am not Get-Process")
+    }
+    Clear-Mocks
+}
+
+Describe "When calling Mock on cmdlet Used by Mock" {
+    Mock Invoke-Command {return "I am not Invoke-Command"}
+    $result = Invoke-Command {return "yes I am"}
+
+    It "Should Invoke the mocked script" {
+        $result.should.be("I am not Invoke-Command")
+    }
+    Clear-Mocks
+}
+
+
 Describe "When calling Mock on non-existing function" {
     try{
         Mock NotFunctionUnderTest {return}
@@ -29,6 +52,29 @@ Describe "When calling Mock on non-existing function" {
     }
 
     It "Should throw correct error" {
-        $result.Exception.Message.should.be("Could not find function NotFunctionUnderTest")
+        $result.Exception.Message.should.be("Could not find command NotFunctionUnderTest")
+    }
+    Clear-Mocks
+}
+
+Describe "When calling Mock on existing function without matching params" {
+    Mock FunctionUnderTest {return "fake results"} -parameterFilter {$param1 -eq "test"}
+
+    $result=FunctionUnderTest "badTest"
+
+    It "Should redirect to real function" {
+        $result.should.be("I am a real world test")
+    }
+    Clear-Mocks
+}
+
+Describe "When calling Mock on existing function with matching params" {
+    Mock FunctionUnderTest {return "fake results"} -parameterFilter {$param1 -eq "badTest"}
+
+    $result=FunctionUnderTest "badTest"
+
+    It "Should return mocked result" {
+        $result.should.be("fake results")
     }
+    Clear-Mocks
 }
diff --git a/Functions/Mock.ps1 b/Functions/Mock.ps1
index c0737cc59..a78de7240 100644
--- a/Functions/Mock.ps1
+++ b/Functions/Mock.ps1
@@ -1,10 +1,24 @@
-function Mock ([string]$function, [ScriptBlock]$mockWith, [switch]$verifiable, [HashTable]$parameterFilters = @{})
+$mockTable = @{}
+
+function Mock ([string]$commandName, [ScriptBlock]$mockWith, [switch]$verifiable, [ScriptBlock]$parameterFilter = {$True})
 {
     # If verifiable, add to a verifiable hashtable
-    if(!(Test-Path Function:\$function)){ Throw "Could not find function $function"}
-    Rename-Item Function:\$function script:PesterIsMocking_$function
-    Set-Item Function:\script:$function -value $mockWith
-    # Mocked function should redirect to real function if param filters are not met
+    $origCommand = (Get-Command $commandName -ErrorAction SilentlyContinue)
+    if(!$origCommand){ Throw "Could not find Command $commandName"}
+    $blocks = @{Mock=$mockWith; Filter=$parameterFilter}
+    $mock = $mockTable.$commandName
+    if(!$mock) {
+        if($origCommand.CommandType -eq "Function") {
+            Rename-Item Function:\$commandName script:PesterIsMocking_$commandName
+        }
+        $metadata=New-Object System.Management.Automation.CommandMetaData $origCommand
+        $cmdLetBinding = [Management.Automation.ProxyCommand]::GetCmdletBindingAttribute($metadata)
+        $params = [Management.Automation.ProxyCommand]::GetParamBlock($metadata)
+        $newContent=Get-Content function:\MockPrototype
+        Set-Item Function:\script:$commandName -value "$cmdLetBinding `r`n param ( $params ) `r`n$newContent"
+        $mock=@{OriginalCommand=$origCommand;blocks=@($blocks)}
+    } else {$mock.blocks += $blocks}
+    $mockTable.$commandName = $mock
     # param filters are met, mark in the verifiable table
 }
 
@@ -16,5 +30,20 @@ function Assert-VerifiableMocks {
 function Clear-Mocks {
     # Called at the end of Describe
     # Clears the Verifiable table
-    # Renames all renamed mocks back to original names
+    $mockTable.Keys | % { Microsoft.PowerShell.Management\Remove-Item function:\$_ }
+    $script:mockTable = @{}
+    Microsoft.PowerShell.Management\Get-ChildItem Function: | ? { $_.Name.StartsWith("PesterIsMocking_") } | % {Microsoft.PowerShell.Management\Rename-Item Function:\$_ "script:$($_.Name.Replace('PesterIsMocking_', ''))"}
+}
+
+function MockPrototype {
+    $functionName = $MyInvocation.MyCommand.Name
+    $mock=$mockTable.$functionName
+    $idx=$mock.blocks.Length
+    while(--$idx -ge 0) {
+        if(Microsoft.PowerShell.Core\Invoke-Command $mock.blocks[$idx].Filter) { 
+            Microsoft.PowerShell.Core\Invoke-Command $mockTable.$functionName.blocks.mock
+            return
+        }
+    }
+    &($mock.OriginalCommand)
 }
\ No newline at end of file