If we assign nil to the association, we should not fail

Author: Kyle d'Oliveira

lib/polymorphic_integer_type/extensions.rb

@@ -16,9 +16,12 @@ def belongs_to(name, options = {})
             mapping[t]
           end
 
-          define_method "#{foreign_type}=" do |klass|
-            klass = klass.base_class if klass.kind_of?(Class)
+          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
             super(enum)
           end
 

lib/polymorphic_integer_type/version.rb

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

spec/polymorphic_integer_type_spec.rb

@@ -17,6 +17,16 @@
 
   let(:link) { Link.create(:source => source, :target => target) }
 
+  context "when the source is nil" do
+    let(:source) { nil }
+    let(:target) { nil }
+    it "should have the no id/type for the source" do
+      expect(link.source_id).to be_nil
+      expect(link.source_type).to be_nil
+      expect(link.source).to be_nil
+    end
+  end
+
   shared_examples "proper source" do
     it "should have the proper id, type and object for the source" do
       expect(link.source_id).to eql source.id