Fix any_instance so that it updates already stubbed instances

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.rb

@@ -67,10 +67,7 @@ def self.teardown
     #   x = 0
     #   RSpec::Mocks.allow_message(bar, :foo) { x += 1 }
     def self.allow_message(subject, message, opts={}, &block)
-      orig_caller = opts.fetch(:expected_from) {
-        CallerFilter.first_non_rspec_line
-      }
-      space.proxy_for(subject).add_stub(orig_caller, message, opts, &block)
+      space.proxy_for(subject).add_stub(message, opts, &block)
     end
 
     # Sets a message expectation on `subject`.
@@ -85,10 +82,7 @@ def self.allow_message(subject, message, opts={}, &block)
     #   RSpec::Mocks.expect_message(bar, :foo)
     #   bar.foo
     def self.expect_message(subject, message, opts={}, &block)
-      orig_caller = opts.fetch(:expected_from) {
-        CallerFilter.first_non_rspec_line
-      }
-      space.proxy_for(subject).add_message_expectation(orig_caller, message, opts, &block)
+      space.proxy_for(subject).add_message_expectation(message, opts, &block)
     end
 
     # Call the passed block and verify mocks after it has executed. This allows

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/expectation_chain.rb

@@ -25,8 +25,10 @@ class PositiveExpectationChain < ExpectationChain
 
         def create_message_expectation_on(instance)
           proxy = ::RSpec::Mocks.space.proxy_for(instance)
-          expected_from = IGNORED_BACKTRACE_LINE
-          me = proxy.add_message_expectation(expected_from, *@expectation_args, &@expectation_block)
+          method_name, opts = @expectation_args
+          opts = (opts || {}).merge(:expected_form => IGNORED_BACKTRACE_LINE)
+
+          me = proxy.add_message_expectation(method_name, opts, &@expectation_block)
           if RSpec::Mocks.configuration.yield_receiver_to_any_instance_implementation_blocks?
             me.and_yield_receiver_to_implementation
           end

lib/rspec/mocks/any_instance/proxy.rb

@@ -0,0 +1,112 @@
+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
+
+        if RUBY_VERSION.to_f > 1.8
+          def respond_to_missing?(method_name, include_private = false)
+            super || @targets.first.respond_to?(method_name, include_private)
+          end
+        else
+          def respond_to?(method_name, include_private = false)
+            super || @targets.first.respond_to?(method_name, include_private)
+          end
+        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/any_instance/stub_chain.rb

@@ -13,8 +13,10 @@ def expectation_fulfilled?
 
         def create_message_expectation_on(instance)
           proxy = ::RSpec::Mocks.space.proxy_for(instance)
-          expected_from = IGNORED_BACKTRACE_LINE
-          stub = proxy.add_stub(expected_from, *@expectation_args, &@expectation_block)
+          method_name, opts = @expectation_args
+          opts = (opts || {}).merge(:expected_form => IGNORED_BACKTRACE_LINE)
+
+          stub = proxy.add_stub(method_name, opts, &@expectation_block)
           @recorder.stubs[stub.message] << stub
 
           if RSpec::Mocks.configuration.yield_receiver_to_any_instance_implementation_blocks?

lib/rspec/mocks/matchers/receive.rb

@@ -9,7 +9,6 @@ def initialize(message, block)
           @message                 = message
           @block                   = block
           @recorded_customizations = []
-          @backtrace_line          = CallerFilter.first_non_rspec_line
         end
 
         def name
@@ -62,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 " <<
@@ -74,12 +73,12 @@ def warn_if_any_instance(expression, subject)
 
         def setup_mock_proxy_method_substitute(subject, method, block)
           proxy = ::RSpec::Mocks.space.proxy_for(subject)
-          setup_method_substitute(proxy, method, block, @backtrace_line)
+          setup_method_substitute(proxy, 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

@@ -42,7 +42,8 @@ def original_method_handle_for(message)
       end
 
       # @private
-      def add_message_expectation(location, method_name, opts={}, &block)
+      def add_message_expectation(method_name, opts={}, &block)
+        location = opts.fetch(:expected_from) { CallerFilter.first_non_rspec_line }
         meth_double = method_double_for(method_name)
 
         if null_object? && !block
@@ -101,7 +102,8 @@ def check_for_unexpected_arguments(expectation)
       end
 
       # @private
-      def add_stub(location, method_name, opts={}, &implementation)
+      def add_stub(method_name, opts={}, &implementation)
+        location = opts.fetch(:expected_from) { CallerFilter.first_non_rspec_line }
         method_double_for(method_name).add_stub @error_generator, @order_group, location, opts, &implementation
       end
 
@@ -116,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
@@ -380,7 +382,7 @@ def initialize(order_group)
       attr_accessor :warn_about_expectations
       alias warn_about_expectations? warn_about_expectations
 
-      def add_message_expectation(location, method_name, opts={}, &block)
+      def add_message_expectation(method_name, opts={}, &block)
         warn(method_name) if warn_about_expectations?
         super
       end
@@ -390,7 +392,7 @@ def add_negative_message_expectation(location, method_name, &implementation)
         super
       end
 
-      def add_stub(location, method_name, opts={}, &implementation)
+      def add_stub(method_name, opts={}, &implementation)
         warn(method_name) if warn_about_expectations?
         super
       end

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

@@ -30,22 +30,19 @@ def self.enable_should(syntax_host = default_should_syntax_host)
         syntax_host.class_exec do
           def should_receive(message, opts={}, &block)
             ::RSpec::Mocks::Syntax.warn_unless_should_configured(__method__)
-            opts[:expected_from] ||= CallerFilter.first_non_rspec_line
             ::RSpec::Mocks.expect_message(self, message, opts, &block)
           end
 
           def should_not_receive(message, &block)
             ::RSpec::Mocks::Syntax.warn_unless_should_configured(__method__)
-            opts = {:expected_from => CallerFilter.first_non_rspec_line}
-            ::RSpec::Mocks.expect_message(self, message, opts, &block).never
+            ::RSpec::Mocks.expect_message(self, message, {}, &block).never
           end
 
           def stub(message_or_hash, opts={}, &block)
             ::RSpec::Mocks::Syntax.warn_unless_should_configured(__method__)
             if ::Hash === message_or_hash
               message_or_hash.each {|message, value| stub(message).and_return value }
             else
-              opts[:expected_from] = CallerFilter.first_non_rspec_line
               ::RSpec::Mocks.allow_message(self, message_or_hash, opts, &block)
             end
           end
@@ -80,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