lib/rr/scenario_creator.rb in rr-0.3.3 vs lib/rr/scenario_creator.rb in rr-0.3.4

@@ -7,38 +7,21 @@
   #
   # Probing can also be added.
   class ScenarioCreator
     NO_SUBJECT_ARG = Object.new
 
-    attr_reader :space, :subject
+    attr_reader :space
     include Errors
 
     def initialize(space)
       @space = space
       @strategy = nil
       @probe = false
       @instance_of = nil
+      @instance_of_method_name = nil
     end
     
-    def create!(subject, method_name, *args, &handler)
-      @subject = subject
-      @method_name = method_name
-      @args = args
-      @handler = handler
-      @double = @space.double(@subject, method_name)
-      @scenario = @space.scenario(@double)
-      transform!
-      @scenario
-    end
-
-#    def instance_of(subject=NO_SUBJECT_ARG, method_name=nil, &definition)
-#      return self if subject === NO_SUBJECT_ARG
-#      raise ArgumentError, "instance_of only accepts class objects" unless subject.is_a?(Class)
-#      @instance_of = true
-#      RR::Space.scenario_method_proxy(self, subject, method_name, &definition)
-#    end
-
     # This method sets the Scenario to have a mock strategy. A mock strategy
     # sets the default state of the Scenario to expect the method call
     # with arguments exactly one time. The Scenario's expectations can be
     # changed.
     #
@@ -182,56 +165,84 @@
       probe_when_do_not_call_error! if @strategy == :do_not_call
       @probe = true
       return self if subject.__id__ === NO_SUBJECT_ARG.__id__
       RR::Space.scenario_method_proxy(self, subject, method_name, &definition)
     end
-    
-    protected
-    def transform!
-      case @strategy
-      when :mock; mock!
-      when :stub; stub!
-      when :do_not_call; do_not_call!
-      else no_strategy_error!
-      end
-      
-      if @probe
-        probe!
+
+    # Calling instance_of will cause all instances of the passed in Class
+    # to have the Scenario defined.
+    #
+    # The following example mocks all User's valid? method and return false. 
+    #   mock.instance_of(User).valid? {false}
+    #
+    # The following example mocks and probes User#projects and returns the
+    # first 3 projects.
+    #   mock.instance_of(User).projects do |projects|
+    #     projects[0..2]
+    #   end
+    def instance_of(subject=NO_SUBJECT_ARG, method_name=nil, &definition)
+      @instance_of = true
+      return self if subject === NO_SUBJECT_ARG
+      raise ArgumentError, "instance_of only accepts class objects" unless subject.is_a?(Class)
+      RR::Space.scenario_method_proxy(self, subject, method_name, &definition)
+    end
+
+    def create!(subject, method_name, *args, &handler)
+      @args = args
+      @handler = handler
+      if @instance_of
+        setup_class_probing_instances(subject, method_name)
       else
-        reimplementation!
+        setup_scenario(subject, method_name)
       end
+      transform!
+      @definition
     end
-
-    def mock!
-      @scenario.with(*@args).once
+    
+    protected
+    def setup_scenario(subject, method_name)
+      @double = @space.double(subject, method_name)
+      @scenario = @space.scenario(@double)
+      @definition = @scenario.definition
     end
 
-    def stub!
-      @scenario.any_number_of_times
-      permissive_argument!
-    end
+    def setup_class_probing_instances(subject, method_name)
+      class_double = @space.double(subject, :new)
+      class_scenario = @space.scenario(class_double)
 
-    def do_not_call!
-      @scenario.never
-      permissive_argument!
-      reimplementation!
-    end
+      instance_method_name = method_name
 
-    def permissive_argument!
-      if @args.empty?
-        @scenario.with_any_args
-      else
-        @scenario.with(*@args)
+      @definition = @space.scenario_definition
+      class_handler = proc do |return_value|
+        double = @space.double(return_value, instance_method_name)
+        @space.scenario(double, @definition)
+        return_value
       end
-    end
 
-    def reimplementation!
-      @scenario.returns(&@handler)
+      builder = ScenarioDefinitionBuilder.new(
+        class_scenario.definition,
+        [],
+        class_handler
+      )
+      builder.stub!
+      builder.probe!
     end
-    
-    def probe!
-      @scenario.implemented_by_original_method
-      @scenario.after_call(&@handler) if @handler
+
+    def transform!
+      builder = ScenarioDefinitionBuilder.new(@definition, @args, @handler)
+
+      case @strategy
+      when :mock; builder.mock!
+      when :stub; builder.stub!
+      when :do_not_call; builder.do_not_call!
+      else no_strategy_error!
+      end
+      
+      if @probe
+        builder.probe!
+      else
+        builder.reimplementation!
+      end
     end
 
     def strategy_error!
       raise(
         ScenarioDefinitionError,