From 9f2ad80167aeec147f3056bb0bee4d16d18c14a2 Mon Sep 17 00:00:00 2001
From: Derick Bailey <derickbailey@gmail.com>
Date: Sun, 25 Oct 2009 13:05:40 -0500
Subject: [PATCH] added ability to provide the stub method block after calling
 the .yields() method, and execute the block as the stubbed method's code

---
 lib/not_a_mock/stubmethod.rb        |  3 ++-
 spec/stub_method_with_block_spec.rb | 29 +++++++++++++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/lib/not_a_mock/stubmethod.rb b/lib/not_a_mock/stubmethod.rb
index 9b59fa9..a89856f 100644
--- a/lib/not_a_mock/stubmethod.rb
+++ b/lib/not_a_mock/stubmethod.rb
@@ -9,8 +9,9 @@ def initialize(&block)
 			@yield_values = []
 		end
 		
-		def yields(*args)
+		def yields(*args, &block)
 			@yield_values = args
+			(@block = block) unless block.nil?
 		end
 		
 		def execute_return_block(*args)
diff --git a/spec/stub_method_with_block_spec.rb b/spec/stub_method_with_block_spec.rb
index bf3787c..da786b8 100644
--- a/spec/stub_method_with_block_spec.rb
+++ b/spec/stub_method_with_block_spec.rb
@@ -78,6 +78,10 @@ def do_something
 		}
 	end
 	
+	after :each do
+		NotAMock::Stubber.instance.reset
+	end
+	
 	it "should yield the first value" do
 		@yielded_bool.should be_true
 	end
@@ -101,6 +105,10 @@ def do_something
 		@yieldtest = YieldTest.new
 		@yieldtest.do_something
 	end
+
+	after :each do
+		NotAMock::Stubber.instance.reset
+	end
 	
 	it "should track the method call on the yielded object" do
 		@stubyielded.should have_received(:yielding_test)
@@ -109,4 +117,25 @@ def do_something
 	it "should track the method call on the yielding object" do
 		YieldingObject.should have_received(:yield_test)
 	end
+end
+
+describe NotAMock::Stubber, "when stubbing a method and providing a block in the yields call" do
+	before :each do
+		@bt = BlockTest.new
+
+		@block_was_called = false
+		@output = @bt.stub_method(:method_that_yields_a_value).yields(true){
+			@block_was_called = true
+		} 
+		
+		@bt.method_that_yields_a_value()
+	end	
+	
+	after :each do
+		NotAMock::Stubber.instance.reset
+	end
+	
+	it "should execute the block provided to the yield, as the stub method block" do
+		@block_was_called.should be_true
+	end
 end
\ No newline at end of file