Check the subclass first before checking the base_class of the sti

Author: Kyle d'Oliveira

lib/polymorphic_integer_type/extensions.rb

@@ -17,11 +17,9 @@ def belongs_to(name, options = {})
           end
 
           define_method "#{foreign_type}=" do |klass|            
-            klass = klass.base_class if klass.kind_of?(Class) && klass <= ActiveRecord::Base
-            enum = mapping.key(klass.to_s) || klass
-            if enum == NilClass
-              enum = nil
-            end
+            enum = mapping.key(klass.to_s)
+            enum ||= mapping.key(klass.base_class.to_s) if klass.kind_of?(Class) && klass <= ActiveRecord::Base
+            enum ||= klass if klass != NilClass
             super(enum)
           end
 

lib/polymorphic_integer_type/version.rb

@@ -1,3 +1,3 @@
 module PolymorphicIntegerType
-  VERSION = "1.0.4"
+  VERSION = "1.0.5"
 end