Make any_instance update existing instance stubs.

Fixes #613.

Changelog.md

@@ -37,6 +37,8 @@ Bug Fixes:
   on an any instance. (Xavier Shay)
 * Fix `and_call_original` to work properly when multiple classes in an
   inheritance hierarchy have been stubbed with the same method. (Myron Marston)
+* Fix `any_instance` so that it updates existing instances that have
+  already been stubbed. (Myron Marston)
 
 ### 3.0.0.beta2 / 2014-02-17
 [Full Changelog](http://github.com/rspec/rspec-mocks/compare/v3.0.0.beta1...v3.0.0.beta2)

lib/rspec/mocks/any_instance.rb

@@ -6,4 +6,5 @@
   any_instance/expectation_chain
   any_instance/message_chains
   any_instance/recorder
+  any_instance/proxy
 ].each { |f| RSpec::Support.require_rspec_mocks(f) }

lib/rspec/mocks/any_instance/proxy.rb

@@ -0,0 +1,102 @@
+module RSpec
+  module Mocks
+    module AnyInstance
+      # @private
+      # The `AnyInstance::Recorder` is responsible for redefining the klass's
+      # instance method in order to add any stubs/expectations the first time
+      # the method is called. It's not capable of updating a stub on an instance
+      # that's already been previously stubbed (either directly, or via
+      # `any_instance`).
+      #
+      # This proxy sits in front of the recorder and delegates both to it
+      # and to the `RSpec::Mocks::Proxy` for each already mocked or stubbed
+      # instance of the class, in order to propogates changes to the instances.
+      #
+      # Note that unlike `RSpec::Mocks::Proxy`, this proxy class is stateless
+      # and is not persisted in `RSpec::Mocks.space`.
+      #
+      # Proxying for the message expectation fluent interface (typically chained
+      # off of the return value of one of these methods) is provided by the
+      # `FluentInterfaceProxy` class below.
+      class Proxy
+        def initialize(recorder, target_proxies)
+          @recorder       = recorder
+          @target_proxies = target_proxies
+        end
+
+        def klass
+          @recorder.klass
+        end
+
+        def stub(method_name_or_method_map, &block)
+          if Hash === method_name_or_method_map
+            method_name_or_method_map.each do |method_name, return_value|
+              stub(method_name).and_return(return_value)
+            end
+          else
+            perform_proxying(__method__, [method_name_or_method_map], block) do |proxy|
+              proxy.add_stub(method_name_or_method_map, &block)
+            end
+          end
+        end
+
+        def unstub(method_name)
+          perform_proxying(__method__, [method_name], nil) do |proxy|
+            proxy.remove_stub_if_present(method_name)
+          end
+        end
+
+        def stub_chain(*chain, &block)
+          perform_proxying(__method__, chain, block) do |proxy|
+            Mocks::StubChain.stub_chain_on(proxy.object, *chain, &block)
+          end
+        end
+
+        def expect_chain(*chain, &block)
+          perform_proxying(__method__, chain, block) do |proxy|
+            Mocks::ExpectChain.expect_chain_on(proxy.object, *chain, &block)
+          end
+        end
+
+        def should_receive(method_name, &block)
+          perform_proxying(__method__, [method_name], block) do |proxy|
+            proxy.add_message_expectation(method_name, &block)
+          end
+        end
+
+        def should_not_receive(method_name, &block)
+          perform_proxying(__method__, [method_name], block) do |proxy|
+            proxy.add_message_expectation(method_name, &block).never
+          end
+        end
+
+      private
+
+        def perform_proxying(method_name, args, block, &target_proxy_block)
+          recorder_value = @recorder.__send__(method_name, *args, &block)
+          proxy_values   = @target_proxies.map(&target_proxy_block)
+          FluentInterfaceProxy.new([recorder_value] + proxy_values)
+        end
+      end
+
+      # @private
+      # Delegates messages to each of the given targets in order to
+      # provide the fluent interface that is available off of message
+      # expectations when dealing with `any_instance`.
+      #
+      # `targets` will typically contain 1 of the `AnyInstance::Recorder`
+      # return values and N `MessageExpectation` instances (one per instance
+      # of the `any_instance` klass).
+      class FluentInterfaceProxy
+        def initialize(targets)
+          @targets = targets
+        end
+
+        def method_missing(*args, &block)
+          return_values = @targets.map { |t| t.__send__(*args, &block) }
+          FluentInterfaceProxy.new(return_values)
+        end
+      end
+    end
+  end
+end

lib/rspec/mocks/any_instance/recorder.rb

@@ -26,15 +26,9 @@ def initialize(klass)
         # instance of this object that invokes the submitted method.
         #
         # @see Methods#stub
-        def stub(method_name_or_method_map, &block)
-          if Hash === method_name_or_method_map
-            method_name_or_method_map.each do |method_name, return_value|
-              stub(method_name).and_return(return_value)
-            end
-          else
-            observe!(method_name_or_method_map)
-            message_chains.add(method_name_or_method_map, StubChain.new(self, method_name_or_method_map, &block))
-          end
+        def stub(method_name, &block)
+          observe!(method_name)
+          message_chains.add(method_name, StubChain.new(self, method_name, &block))
         end
 
         # Initializes the recording a stub chain to be played back against any
@@ -85,9 +79,6 @@ 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.space.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/matchers/receive.rb

@@ -61,7 +61,7 @@ def setup_any_instance_allowance(subject, &block)
       private
 
         def warn_if_any_instance(expression, subject)
-          if AnyInstance::Recorder === subject
+          if AnyInstance::Proxy === subject
             RSpec.warning(
               "`#{expression}(#{subject.klass}.any_instance).to` " <<
               "is probably not what you meant, it does not operate on " <<
@@ -77,8 +77,8 @@ def setup_mock_proxy_method_substitute(subject, method, block)
         end
 
         def setup_any_instance_method_substitute(subject, method, block)
-          any_instance_recorder = ::RSpec::Mocks.space.any_instance_recorder_for(subject)
-          setup_method_substitute(any_instance_recorder, method, block)
+          proxy = ::RSpec::Mocks.space.any_instance_proxy_for(subject)
+          setup_method_substitute(proxy, method, block)
         end
 
         def setup_method_substitute(host, method, block, *args)

lib/rspec/mocks/matchers/receive_message_chain.rb

@@ -28,14 +28,14 @@ def setup_allowance(subject, &block)
         end
 
         def setup_any_instance_allowance(subject, &block)
-          recorder = ::RSpec::Mocks.space.any_instance_recorder_for(subject)
-          chain = recorder.stub_chain(*@chain, &(@block || block))
+          proxy = ::RSpec::Mocks.space.any_instance_proxy_for(subject)
+          chain = proxy.stub_chain(*@chain, &(@block || block))
           replay_customizations(chain)
         end
 
         def setup_any_instance_expectation(subject, &block)
-          recorder = ::RSpec::Mocks.space.any_instance_recorder_for(subject)
-          chain = recorder.expect_chain(*@chain, &(@block || block))
+          proxy = ::RSpec::Mocks.space.any_instance_proxy_for(subject)
+          chain = proxy.expect_chain(*@chain, &(@block || block))
           replay_customizations(chain)
         end
 

lib/rspec/mocks/matchers/receive_messages.rb

@@ -58,7 +58,7 @@ def proxy_on(subject)
         end
 
         def any_instance_of(subject)
-          ::RSpec::Mocks.space.any_instance_recorder_for(subject)
+          ::RSpec::Mocks.space.any_instance_proxy_for(subject)
         end
 
         def each_message_on(host)

lib/rspec/mocks/method_double.rb

@@ -189,13 +189,12 @@ def add_default_stub(*args, &implementation)
       # @private
       def remove_stub
         raise_method_not_stubbed_error if stubs.empty?
-        expectations.empty? ? reset : stubs.clear
+        remove_stub_if_present
       end
 
       # @private
-      def remove_single_stub(stub)
-        stubs.delete(stub)
-        restore_original_method if stubs.empty? && expectations.empty?
+      def remove_stub_if_present
+        expectations.empty? ? reset : stubs.clear
       end
 
       # @private

lib/rspec/mocks/proxy.rb

@@ -118,8 +118,8 @@ def remove_stub(method_name)
       end
 
       # @private
-      def remove_single_stub(method_name, stub)
-        method_double_for(method_name).remove_single_stub(stub)
+      def remove_stub_if_present(method_name)
+        method_double_for(method_name).remove_stub_if_present
       end
 
       # @private

lib/rspec/mocks/space.rb

@@ -13,6 +13,10 @@ def any_instance_recorder_for(*args)
         raise_lifecycle_message
       end
 
+      def any_instance_proxy_for(*args)
+        raise_lifecycle_message
+      end
+
       def register_constant_mutator(mutator)
         raise_lifecycle_message
       end
@@ -85,6 +89,10 @@ def any_instance_recorder_for(klass)
         end
       end
 
+      def any_instance_proxy_for(klass)
+        AnyInstance::Proxy.new(any_instance_recorder_for(klass), proxies_of(klass))
+      end
+
       def proxies_of(klass)
         proxies.values.select { |proxy| klass === proxy.object }
       end

lib/rspec/mocks/syntax.rb

@@ -77,7 +77,7 @@ def received_message?(message, *args, &block)
             Class.class_exec do
               def any_instance
                 ::RSpec::Mocks::Syntax.warn_unless_should_configured(__method__)
-                ::RSpec::Mocks.space.any_instance_recorder_for(self)
+                ::RSpec::Mocks.space.any_instance_proxy_for(self)
               end
             end
           end

spec/rspec/mocks/any_instance_spec.rb

@@ -175,6 +175,8 @@ def private_method; :private_method_return_value; end
             expect(instance.foo).to eq(1)
             allow_any_instance_of(klass).to receive(:foo).and_return(2)
             expect(instance.foo).to eq(2)
+            allow_any_instance_of(klass).to receive(:foo).and_raise("boom")
+            expect { instance.foo }.to raise_error("boom")
           end
         end
 
@@ -308,7 +310,6 @@ class RSpec::SampleRspecTestClass;end
           obj.existing_method
           allow_any_instance_of(klass).to receive(:existing_method).and_call_original
 
-          pending "not working for `and_call_original` yet, but works with `unstub`"
           expect(obj.existing_method).to eq(:existing_method_return_value)
         end
 
@@ -318,16 +319,15 @@ class RSpec::SampleRspecTestClass;end
           expect(obj.existing_method).to eq(:any_instance_value)
           allow_any_instance_of(klass).to receive(:existing_method).and_call_original
 
-          pending "not working for `and_call_original` yet, but works with `unstub`"
           expect(obj.existing_method).to eq(:existing_method_return_value)
         end
 
-        it "does not remove any stubs set directly on an instance" do
+        it "removes any stubs set directly on an instance" do
           allow_any_instance_of(klass).to receive(:existing_method).and_return(:any_instance_value)
           obj = klass.new
           allow(obj).to receive(:existing_method).and_return(:local_method)
           allow_any_instance_of(klass).to receive(:existing_method).and_call_original
-          expect(obj.existing_method).to eq(:local_method)
+          expect(obj.existing_method).to eq(:existing_method_return_value)
         end
 
         it "does not remove any expectations with the same method name" do
@@ -360,6 +360,15 @@ class RSpec::SampleRspecTestClass;end
           expect { klass.new.another_existing_method }.to_not raise_error
         end
 
+        it "affects previously stubbed instances" do
+          instance = klass.new
+
+          allow_any_instance_of(klass).to receive(:foo).and_return(1)
+          expect(instance.foo).to eq(1)
+          expect_any_instance_of(klass).not_to receive(:foo)
+          expect { instance.foo }.to fail
+        end
+
         context "with constraints" do
           it "fails if the method is called with the specified parameters" do
             expect_any_instance_of(klass).not_to receive(:existing_method_with_arguments).with(:argument_one, :argument_two)
@@ -409,6 +418,18 @@ def inspect
           reset_all
         end
 
+        it "affects previously stubbed instances" do
+          instance = klass.new
+
+          allow_any_instance_of(klass).to receive(:foo).and_return(1)
+          expect(instance.foo).to eq(1)
+          expect_any_instance_of(klass).to receive(:foo).with(2).and_return(2)
+          expect(instance.foo(2)).to eq(2)
+
+          # TODO: this shouldn't be necessary to satisfy the expectation, but is.
+          klass.new.foo(2)
+        end
+
         context "with an expectation is set on a method which does not exist" do
           it "returns the expected value" do
             expect_any_instance_of(klass).to receive(:foo).and_return(1)

spec/rspec/mocks/matchers/receive_message_chain_spec.rb

@@ -143,13 +143,33 @@ module RSpec::Mocks::Matchers
         expect(o.to_a.length).to eq(3)
       end
 
+      it "stubs already stubbed instances when using `allow_any_instance_of`" do
+        o = Object.new
+        allow(o).to receive(:foo).and_return(dbl = double)
+        expect(o.foo).to be(dbl)
+
+        allow_any_instance_of(Object).to receive_message_chain(:foo, :bar).and_return("bazz")
+        expect(o.foo.bar).to eq("bazz")
+      end
+
       it "fails when with expect_any_instance_of is used and the entire chain is not called" do
         expect {
           expect_any_instance_of(Object).to receive_message_chain(:to_a, :length => 3)
           verify_all
         }.to raise_error(RSpec::Mocks::MockExpectationError)
       end
 
+      it "affects previously stubbed instances when `expect_any_instance_of` is called" do
+        o = Object.new
+        allow(o).to receive(:foo).and_return(double)
+
+        expect_any_instance_of(Object).to receive_message_chain(:foo, :bar => 3)
+        expect(o.foo.bar).to eq(3)
+
+        # TODO: this shouldn't be necessary to satisfy the expectation, but is.
+        Object.new.foo.bar
+      end
+
       it "passes when with expect_any_instance_of is used and the entire chain is called" do
         o = Object.new
 

spec/rspec/mocks/matchers/receive_messages_spec.rb

@@ -52,6 +52,15 @@ module Mocks
         expect(obj.b).to eq 2
       end
 
+      it "updates stubs on instances with existing stubs" do
+        allow(obj).to receive(:a).and_return(3)
+        expect(obj.a).to eq(3)
+
+        allow_any_instance_of(Object).to receive_messages(:a => 1, :b => 2)
+        expect(obj.a).to eq 1
+        expect(obj.b).to eq 2
+      end
+
       it_behaves_like "complains when given blocks"
     end
 

spec/rspec/mocks/matchers/receive_spec.rb

@@ -18,14 +18,16 @@ module Mocks
 
         it "warns about expect(Klass.any_instance).to receive..." do
           expect(RSpec).to receive(:warning).with(/expect.*any_instance.*is probably not what you meant.*expect_any_instance_of.*instead/)
-          expect(Object.any_instance).to receive(:foo)
-          Object.any_instance.foo
+          any_instance_proxy = Object.any_instance
+          expect(any_instance_proxy).to receive(:foo)
+          any_instance_proxy.foo
         end
 
         it "includes the correct call site in the expect warning" do
+          any_instance_proxy = Object.any_instance
           expect_warning_with_call_site(__FILE__, __LINE__ + 1)
-          expect(Object.any_instance).to receive(:foo)
-          Object.any_instance.foo
+          expect(any_instance_proxy).to receive(:foo)
+          any_instance_proxy.foo
         end
       end