Fix unstubing existing instances of any instance

* access proxies for any instance of klass
* remove stubs from "running" instances
* pass recorder into chains to register stubs
* reduce lookup cost of proxies by klass
* record stubs used
* remove just the stubs recorded by any instance
* switch from for..in to each
* allow unstubing of instances that are a sub class
* refactor lookup for subclass
* Improve doc string.
* "local instance" didn't really give me the right
* sense for what this was testing.
* Improve space spec.
    - The old spec only checked the number of proxies returned,
      and didn't actually check that it returned the right ones.
    - The old spec only tested who were instances of the given
      class, and did not check instances of subclasses.

Changelog.md

@@ -35,6 +35,8 @@ Bug Fixes:
 
 * Fix `and_call_original` to handle a complex edge case involving
   singleton class ancestors. (Marc-André Lafortune, Myron Marston)
+* Fix issue where unstubing methods on "any instances" would not
+  remove stubs on existing instances (Jon Rowe)
 
 ### 2.14.3 / 2013-08-08
 [full changelog](http://github.com/rspec/rspec-mocks/compare/v2.14.2...v2.14.3)

lib/rspec/mocks.rb

@@ -26,6 +26,10 @@ def proxy_for(object)
         space.proxy_for(object)
       end
 
+      def proxies_of(klass)
+        space.proxies_of(klass)
+      end
+
       def any_instance_recorder_for(klass)
         space.any_instance_recorder_for(klass)
       end

lib/rspec/mocks/any_instance/chain.rb

@@ -2,7 +2,8 @@ module RSpec
   module Mocks
     module AnyInstance
       class Chain
-        def initialize(*args, &block)
+        def initialize(recorder, *args, &block)
+          @recorder          = recorder
           @expectation_args  = args
           @expectation_block = block
         end

lib/rspec/mocks/any_instance/recorder.rb

@@ -11,10 +11,11 @@ module AnyInstance
       # @see Chain
       class Recorder
         # @private
-        attr_reader :message_chains
+        attr_reader :message_chains, :stubs
 
         def initialize(klass)
           @message_chains = MessageChains.new
+          @stubs = Hash.new { |hash,key| hash[key] = [] }
           @observed_methods = []
           @played_methods = {}
           @klass = klass
@@ -32,7 +33,7 @@ def stub(method_name_or_method_map, &block)
             end
           else
             observe!(method_name_or_method_map)
-            message_chains.add(method_name_or_method_map, StubChain.new(method_name_or_method_map, &block))
+            message_chains.add(method_name_or_method_map, StubChain.new(self, method_name_or_method_map, &block))
           end
         end
 
@@ -44,7 +45,7 @@ def stub(method_name_or_method_map, &block)
         def stub_chain(*method_names_and_optional_return_values, &block)
           normalize_chain(*method_names_and_optional_return_values) do |method_name, args|
             observe!(method_name)
-            message_chains.add(method_name, StubChainChain.new(*args, &block))
+            message_chains.add(method_name, StubChainChain.new(self, *args, &block))
           end
         end
 
@@ -56,7 +57,7 @@ def stub_chain(*method_names_and_optional_return_values, &block)
         def should_receive(method_name, &block)
           @expectation_set = true
           observe!(method_name)
-          message_chains.add(method_name, PositiveExpectationChain.new(method_name, &block))
+          message_chains.add(method_name, PositiveExpectationChain.new(self, method_name, &block))
         end
 
         def should_not_receive(method_name, &block)
@@ -72,6 +73,10 @@ def unstub(method_name)
             raise RSpec::Mocks::MockExpectationError, "The method `#{method_name}` was not stubbed or was already unstubbed"
           end
           message_chains.remove_stub_chains_for!(method_name)
+          ::RSpec::Mocks.proxies_of(@klass).each do |proxy|
+            stubs[method_name].each { |stub| proxy.remove_single_stub(method_name, stub) }
+          end
+          stubs[method_name].clear
           stop_observing!(method_name) unless message_chains.has_expectation?(method_name)
         end
 

lib/rspec/mocks/any_instance/stub_chain.rb

@@ -15,6 +15,7 @@ def create_message_expectation_on(instance)
           proxy = ::RSpec::Mocks.proxy_for(instance)
           expected_from = IGNORED_BACKTRACE_LINE
           stub = proxy.add_stub(expected_from, *@expectation_args, &@expectation_block)
+          @recorder.stubs[stub.message] << stub
 
           if RSpec::Mocks.configuration.yield_receiver_to_any_instance_implementation_blocks?
             stub.and_yield_receiver_to_implementation

lib/rspec/mocks/method_double.rb

@@ -161,6 +161,12 @@ def remove_stub
         expectations.empty? ? reset : stubs.clear
       end
 
+      # @private
+      def remove_single_stub(stub)
+        stubs.delete(stub)
+        restore_original_method if stubs.empty? && expectations.empty?
+      end
+
       # @private
       def raise_method_not_stubbed_error
         raise MockExpectationError, "The method `#{method_name}` was not stubbed or was already unstubbed"

lib/rspec/mocks/proxy.rb

@@ -15,6 +15,9 @@ def initialize(object, name=nil, options={})
         @null_object = false
       end
 
+      # @private
+      attr_reader :object
+
       # @private
       def null_object?
         @null_object
@@ -102,6 +105,7 @@ def add_stub(location, method_name, opts={}, &implementation)
         method_double[method_name].add_stub @error_generator, @expectation_ordering, location, opts, &implementation
       end
 
+      # @private
       def add_simple_stub(method_name, response)
         method_double[method_name].add_simple_stub method_name, response
       end
@@ -111,6 +115,11 @@ def remove_stub(method_name)
         method_double[method_name].remove_stub
       end
 
+      # @private
+      def remove_single_stub(method_name, stub)
+        method_double[method_name].remove_single_stub(stub)
+      end
+
       # @private
       def verify
         method_doubles.each {|d| d.verify}

lib/rspec/mocks/space.rb

@@ -5,8 +5,8 @@ class Space
       attr_reader :proxies, :any_instance_recorders
 
       def initialize
-        @proxies = {}
-        @any_instance_recorders = {}
+        @proxies                 = {}
+        @any_instance_recorders  = {}
       end
 
       def verify_all
@@ -46,6 +46,10 @@ def remove_any_instance_recorder_for(klass)
         any_instance_recorders.delete(klass.__id__)
       end
 
+      def proxies_of(klass)
+        proxies.values.select { |proxy| klass === proxy.object }
+      end
+
       def proxy_for(object)
         id = id_for(object)
         proxies.fetch(id) do

spec/rspec/mocks/any_instance/message_chains_spec.rb

@@ -1,9 +1,10 @@
 require 'spec_helper'
 
 describe RSpec::Mocks::AnyInstance::MessageChains do
+  let(:recorder) { double }
   let(:chains) { RSpec::Mocks::AnyInstance::MessageChains.new }
-  let(:stub_chain) { RSpec::Mocks::AnyInstance::StubChain.new }
-  let(:expectation_chain) { RSpec::Mocks::AnyInstance::PositiveExpectationChain.new }
+  let(:stub_chain) { RSpec::Mocks::AnyInstance::StubChain.new recorder }
+  let(:expectation_chain) { RSpec::Mocks::AnyInstance::PositiveExpectationChain.new recorder }
 
   it "knows if a method does not have an expectation set on it" do
     chains.add(:method_name, stub_chain)
@@ -19,7 +20,7 @@
   it "can remove all stub chains" do
     chains.add(:method_name, stub_chain)
     chains.add(:method_name, expectation_chain)
-    chains.add(:method_name, RSpec::Mocks::AnyInstance::StubChain.new)
+    chains.add(:method_name, RSpec::Mocks::AnyInstance::StubChain.new(recorder))
 
     chains.remove_stub_chains_for!(:method_name)
     expect(chains[:method_name]).to eq([expectation_chain])
@@ -33,7 +34,7 @@
 
     it "allows multiple stub chains for a method" do
       chains.add(:method_name, stub_chain)
-      chains.add(:method_name, another_stub_chain = RSpec::Mocks::AnyInstance::StubChain.new)
+      chains.add(:method_name, another_stub_chain = RSpec::Mocks::AnyInstance::StubChain.new(recorder))
       expect(chains[:method_name]).to eq([stub_chain, another_stub_chain])
     end
   end

spec/rspec/mocks/any_instance_spec.rb

@@ -279,6 +279,30 @@ class RSpec::SampleRspecTestClass;end
           expect(klass.new.existing_method).to eq(:existing_method_return_value)
         end
 
+        it "removes stubs even if they have already been invoked" do
+          klass.any_instance.stub(:existing_method).and_return(:any_instance_value)
+          obj = klass.new
+          obj.existing_method
+          klass.any_instance.unstub(:existing_method)
+          expect(obj.existing_method).to eq(:existing_method_return_value)
+        end
+
+        it "removes stubs from sub class after invokation when super class was originally stubbed" do
+          klass.any_instance.stub(:existing_method).and_return(:any_instance_value)
+          obj = Class.new(klass).new
+          expect(obj.existing_method).to eq(:any_instance_value)
+          klass.any_instance.unstub(:existing_method)
+          expect(obj.existing_method).to eq(:existing_method_return_value)
+        end
+
+        it "does not remove any stubs set directly on an instance" do
+          klass.any_instance.stub(:existing_method).and_return(:any_instance_value)
+          obj = klass.new
+          obj.stub(:existing_method).and_return(:local_method)
+          klass.any_instance.unstub(:existing_method)
+          expect(obj.existing_method).to eq(:local_method)
+        end
+
         it "does not remove any expectations with the same method name" do
           klass.any_instance.should_receive(:existing_method_with_arguments).with(3).and_return(:three)
           klass.any_instance.stub(:existing_method_with_arguments).with(1)

spec/rspec/mocks/space_spec.rb

@@ -0,0 +1,32 @@
+require 'spec_helper'
+
+module RSpec::Mocks
+  describe Space do
+
+    describe "#proxies_of(klass)" do
+      let(:space) { Space.new }
+
+      it 'returns proxies' do
+        space.proxy_for("")
+        expect(space.proxies_of(String).map(&:class)).to eq([PartialMockProxy])
+      end
+
+      it 'returns only the proxies whose object is an instance of the given class' do
+        grandparent_class = Class.new
+        parent_class      = Class.new(grandparent_class)
+        child_class       = Class.new(parent_class)
+
+        grandparent = grandparent_class.new
+        parent      = parent_class.new
+        child       = child_class.new
+
+        grandparent_proxy = space.proxy_for(grandparent)
+        parent_proxy      = space.proxy_for(parent)
+        child_proxy       = space.proxy_for(child)
+
+        expect(space.proxies_of(parent_class)).to match_array([parent_proxy, child_proxy])
+      end
+    end
+
+  end
+end