cleaned up the stubber by adding a stubmethod object to hold all the …

…data about the stub, including the block to execute

lib/not_a_mock.rb

@@ -11,5 +11,5 @@
 require 'not_a_mock/matchers/times_matcher'
 require 'not_a_mock/object_extensions'
 require 'not_a_mock/stubber'
-require 'not_a_mock/yielder'
+require 'not_a_mock/stubmethod'
 require 'not_a_mock/stub'

lib/not_a_mock/stubber.rb

@@ -9,49 +9,43 @@ class Stubber
     include Singleton
 
     def initialize
-      @stubbed_methods = []
-      @stubbed_blocks = {}
-      @yield_values = {}
+      @stubbed_methods = {}
     end
 
     # Stub +method+ on +object+ to evalutate +block+ and return the result.
     # 
     # You should call Object#stub_method rathing than calling this directly.
     def stub_method(object, method, &block) #:nodoc:
       unless @stubbed_methods.include?([object, method])
-        @stubbed_methods << [object, method]
-        @stubbed_blocks[[object, method]] = block
-        yielder = NotAMock::Yielder.new()
-        @yield_values[[object, method]] = yielder
+      	stubmethod = NotAMock::StubMethod.new(&block)
+        @stubbed_methods[[object, method]] = stubmethod
         add_hook(object, method)
       end
-      yielder
+      stubmethod
     end
 
     # Remove the stubbed +method+ on +object+.
     #
     # You should call Object#unstub_methods rather than calling this directly.
     def unstub_method(object, method) #:nodoc:
       if @stubbed_methods.delete([object, method])
-      	@stubbed_blocks.delete([object, method])
         remove_hook(object, method)
       end
     end
 
     # Removes all stub methods.
     def reset
-      @stubbed_methods.each do |object, method|
+      @stubbed_methods.each do |key, value|
+      	object = key[0]
+      	method = key[1]
         remove_hook(object, method) 
       end
-      @stubbed_methods = []
+      @stubbed_methods = {}
     end
 
-    def get_block(obj, meth)
-    	@stubbed_blocks[[obj, meth]]
-    end
-
-    def get_yielder(obj, meth)
-    	@yield_values[[obj, meth]]
+    #Retrieve the stub method data and code for stubbed +method+ on the +object+
+    def get_stubmethod(object, method)
+    	@stubbed_methods[[object, method]]
     end
 
   private
@@ -63,16 +57,13 @@ def add_hook(object, method)
       end
       object.instance_eval(<<-EOF, __FILE__, __LINE__)
         def #{method}(*args, &block)
-          yielder = Stubber.instance.get_yielder(self, :#{method})
-          if !yielder.yield_values.empty?
-          	yield *yielder.yield_values if block_given?
+          stubmethod = Stubber.instance.get_stubmethod(self, :#{method})
+          if !stubmethod.yield_values.empty?
+          	yield *stubmethod.yield_values if block_given?
           else
             yield if block_given?
           end
-          meth = Stubber.instance.get_block(self, :#{method})
-          return_value = nil
-          return_value = meth.call(*args) unless meth.nil?
-          return_value
+          return stubmethod.execute_return_block(*args)
         end   
       EOF
     end

lib/not_a_mock/stubmethod.rb

@@ -0,0 +1,22 @@
+require 'singleton'
+
+module NotAMock
+	class StubMethod
+		attr_accessor :yield_values, :return_block
+
+		def initialize(&block)
+			@block = block
+			@yield_values = []
+		end
+
+		def yields(*args)
+			@yield_values = args
+		end
+
+		def execute_return_block(*args)
+			return_value = nil
+			(return_value = @block.call(*args)) unless @block.nil?
+			return_value
+		end
+	end
+end

spec/stub_method_with_block_spec.rb

@@ -39,7 +39,7 @@ def do_something
 end
 
 describe NotAMock::Stubber, "when stubbing a method that is called with a block" do
-	before :all do
+	before :each do
 		@bt = BlockTest.new
 
 		@bt.stub_method(:block_method)
@@ -48,6 +48,10 @@ def do_something
 		@bt.block_method(){ @block_executed = true }
 	end
 
+	after :each do
+		NotAMock::Stubber.instance.reset
+	end
+
 	it "should execute the block" do
 		@block_executed.should be_true
 	end
@@ -62,7 +66,7 @@ def do_something
 end
 
 describe NotAMock::Stubber, "when stubbing a method that yields two values" do
-	before :all do
+	before :each do
 		@bt = BlockTest.new
 
 		p = proc { yield false if block_given?}
@@ -88,7 +92,7 @@ def do_something
 end 
 
 describe NotAMock::Stubber, "when stubbing a method on an object that is yielded to a block" do
-	before :all do
+	before :each do
 		@stubyielded = YieldedObject.new
 		@stubyielded.stub_method(:yielding_test)