Properly get the foreign key for STI models

lib/state_machine/audit_trail/active_record.rb

@@ -6,6 +6,6 @@ def log(object, event, from, to, timestamp = Time.now)
   end
 
   def foreign_key_field(object)
-    object.class.name.foreign_key.to_sym
+    object.class.base_class.name.foreign_key.to_sym
   end
 end

spec/spec_helper.rb

@@ -12,6 +12,7 @@
 
 ActiveRecord::Base.connection.create_table(:test_models) do |t|
   t.string :state
+  t.string :type
   t.timestamps
 end
 
@@ -22,7 +23,6 @@
 end
 
 
-
 def create_transition_table(owner_class, state)
   class_name = "#{owner_class.name}#{state.to_s.camelize}Transition"
 
@@ -49,7 +49,6 @@ class TestModelWithMultipleStateMachinesSecondTransition < ActiveRecord::Base
   belongs_to :test_model
 end
 
-
 class TestModel < ActiveRecord::Base
 
   state_machine :state, :initial => :waiting do # log initial state?
@@ -65,6 +64,9 @@ class TestModel < ActiveRecord::Base
   end
 end
 
+class TestModelDescendant < TestModel
+end
+
 class TestModelWithMultipleStateMachines < ActiveRecord::Base
 
   state_machine :first, :initial => :beginning do

spec/state_machine/audit_trail_spec.rb

@@ -79,4 +79,9 @@
       first_transition.created_at.should be_within(10.seconds).of(Time.now.utc)
     end
   end
+
+  it "should properly grab the class name from STI models" do
+    m = TestModelDescendant.create!
+    lambda { m.start! }.should_not raise_error
+  end
 end