lib/rr/stub_creator.rb in rr-0.1.11 vs lib/rr/stub_creator.rb in rr-0.1.12

@@ -10,27 +10,21 @@
   # The StubCreator also supports a block sytnax.
   #
   #    stub(subject) do |m|
   #      m.method_name(arg1, arg2) { return_value }
   #    end
-  class StubCreator
-    instance_methods.each { |m| undef_method m unless m =~ /^__/ }
-    
-    def initialize(space, subject)
-      @space = space
-      @subject = subject
-      yield(self) if block_given?
-    end
-
-    protected
-    def method_missing(method_name, *args, &returns)
-      double = @space.create_double(@subject, method_name)
-      scenario = @space.create_scenario(double)
-      scenario.returns(&returns).any_number_of_times
-      if args.empty?
-        scenario.with_any_args
-      else
-        scenario.with(*args)
+  class StubCreator < Creator
+    module InstanceMethods
+      protected
+      def method_missing(method_name, *args, &returns)
+        double = @space.create_double(@subject, method_name)
+        scenario = @space.create_scenario(double)
+        scenario.returns(&returns).any_number_of_times
+        if args.empty?
+          scenario.with_any_args
+        else
+          scenario.with(*args)
+        end
       end
     end
   end
 end