Have the foreign_type use the base class to deal with STI classes properly

Author: Kyle d'Oliveira

lib/polymorphic_integer_type/extensions.rb

@@ -17,6 +17,7 @@ def belongs_to(name, options = {})
           end
 
           define_method "#{foreign_type}=" do |klass|
+            klass = klass.base_class if klass.kind_of?(Class)
             enum = mapping.key(klass.to_s) || klass
             super(enum)
           end

spec/polymorphic_integer_type_spec.rb

@@ -131,5 +131,15 @@
 
   end
 
+  context "when the association is an STI table" do
+    let(:link) { Link.create(:source => source, :target => whiskey) }
+    let(:source) { Dog.create(:name => "Bela", :kind => "Dog", :owner => owner) }
+    it "should have the proper id, type and object for the source" do
+      expect(link.source_id).to eql source.id
+      expect(link.source_type).to eql "Animal"
+      expect(link.source).to eql source
+    end
+  end
+  
 
 end

spec/spec_helper.rb

@@ -3,6 +3,7 @@
 require 'support/configuration'
 require 'support/link'
 require 'support/animal'
+require 'support/dog'
 require 'support/person'
 require 'support/food'
 require 'support/drink'

spec/support/dog.rb

@@ -0,0 +1,4 @@
+class Dog < Animal
+
+
+end

spec/support/migrations/2_create_animal_table.rb

@@ -3,6 +3,7 @@ class CreateAnimalTable < ActiveRecord::Migration
   def up
     create_table :animals do |t|
       t.string :name
+      t.string :type
       t.string :kind
       t.integer :owner_id
     end