lib/tspec.rb in tspec-0.1.0 vs lib/tspec.rb in tspec-0.2.0

@@ -17,15 +17,36 @@
   def receive(*type, **types)
     self
   end
 end
 
+class UnboundMethod
+  def return(*types)
+    self
+  end
+
+  def receive(*type, **types)
+    self
+  end
+end
+
+class Method
+  def return(*types)
+    self
+  end
+
+  def receive(*type, **types)
+    self
+  end
+end
+
 module TSpec
   @method_return_type_table = {}
   @method_arguments_type_table = {}
   @before_trace = {}
   @type_error_flag = false
+  DEFINE_METHOD_SYMBOLS = %i(method_added singleton_method_added define_method instance_method method)
 
   def self.value_type_check(value, *types)
     types.any? do |type|
       if type.instance_of?(Array)
         return false unless value.instance_of?(Array)
@@ -42,15 +63,20 @@
   end
 
   TracePoint.trace do |tp|
     case tp.event
       when :call
-        if %i(return receive).include?(tp.method_id) && %i(method_added singleton_method_added define_method).include?(@before_trace[:method_id])
-          klass = (@before_trace[:method_id] == :singleton_method_added) ? @before_trace[:self].singleton_class : @before_trace[:self]
+        if %i(return receive).include?(tp.method_id) && DEFINE_METHOD_SYMBOLS.include?(@before_trace[:method_id])
           ctx = tp.binding
-          key = "#{klass}::#{ctx.receiver}"
 
+          if %i(instance_method method).include?(@before_trace[:method_id])
+            key = "#{tp.self.owner}::#{tp.self.name}"
+          else
+            klass = (@before_trace[:method_id] == :singleton_method_added) ? @before_trace[:self].singleton_class : @before_trace[:self]
+            key = "#{klass}::#{ctx.receiver}"
+          end
+
           case tp.method_id
             when :return
               @method_return_type_table[key] = ctx.local_variable_get(:types)
             when :receive
               @method_arguments_type_table[key] = ctx.local_variable_get(:types)
@@ -64,20 +90,17 @@
 
           if types = @method_arguments_type_table[key]
             arguments = tp.binding.eval("method(:#{tp.method_id}).parameters.map(&:last)")
 
             types.each do |name, type|
+              name = arguments.first if name == type.__id__
 
-              if name != type.__id__
-                unless arguments.include?(name)
-                  @type_error_flag = true
-                  raise NotFoundArgumentNameError, "undefined arguments `#{name}' for #{key}"
-                end
-
-                value = tp.binding.local_variable_get(name)
-              else
-                value = tp.binding.local_variable_get(arguments.first)
+              unless arguments.include?(name)
+                @type_error_flag = true
+                raise NotFoundArgumentNameError, "undefined arguments `#{name}' for #{key}"
               end
+
+              value = tp.binding.local_variable_get(name)
 
               unless value_type_check(value, *type)
                 @type_error_flag = true
                 raise ArgumentTypeError, "##{tp.method_id} '#{name}' variable should be #{type.inspect}, but actual '#{value.inspect}' - #{value.class}"
               end