From a0bb1c404c5a0b51d6c41384b78a819537300312 Mon Sep 17 00:00:00 2001
From: Jeff Dean <jeff@zilkey.com>
Date: Tue, 29 Apr 2008 01:58:05 -0400
Subject: [PATCH] make sure the 2 fields are in sync by using
 before_validation_on_create

---
 lib/persistence/active_record_persistence.rb | 15 +++++++++++++++
 spec/unit/active_record_persistence_spec.rb  | 18 ++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/lib/persistence/active_record_persistence.rb b/lib/persistence/active_record_persistence.rb
index dbe5ef1..42b3c79 100644
--- a/lib/persistence/active_record_persistence.rb
+++ b/lib/persistence/active_record_persistence.rb
@@ -33,6 +33,7 @@ module AASM
         base.send(:include, AASM::Persistence::ActiveRecordPersistence::ReadState) unless base.method_defined?(:aasm_read_state)
         base.send(:include, AASM::Persistence::ActiveRecordPersistence::WriteState) unless base.method_defined?(:aasm_write_state)
         base.send(:include, AASM::Persistence::ActiveRecordPersistence::WriteStateWithoutPersistence) unless base.method_defined?(:aasm_write_state_without_persistence)        
+        base.before_validation_on_create :aasm_ensure_initial_state
       end
 
       module ClassMethods
@@ -91,6 +92,20 @@ module AASM
           @current_state = aasm_read_state
         end
 
+        private
+        
+        # Called from before_validation_on_create to ensure
+        # that if there is a nil value in the underlying aasm_state
+        # column, the initial state is used instead
+        #
+        #   foo = Foo.new
+        #   foo.save
+        #   foo.aasm_state # => the initial state
+        #
+        def aasm_ensure_initial_state
+          send("#{self.class.aasm_column}=", self.aasm_current_state.to_s)
+        end
+
       end
 
       module WriteStateWithoutPersistence
diff --git a/spec/unit/active_record_persistence_spec.rb b/spec/unit/active_record_persistence_spec.rb
index d3c61aa..25e3f10 100644
--- a/spec/unit/active_record_persistence_spec.rb
+++ b/spec/unit/active_record_persistence_spec.rb
@@ -157,6 +157,24 @@ begin
       foo.aasm_current_state.should be_nil
     end
 
+    it "should have aasm_ensure_initial_state" do
+      foo = Foo.new
+      foo.send :aasm_ensure_initial_state
+    end
+
+    it "should call aasm_ensure_initial_state on validation before create" do
+      foo = Foo.new
+      foo.should_receive(:aasm_ensure_initial_state).and_return(true)
+      foo.valid?
+    end
+
+    it "should call aasm_ensure_initial_state on validation before create" do
+      foo = Foo.new
+      foo.stub!(:new_record?).and_return(false)
+      foo.should_not_receive(:aasm_ensure_initial_state)
+      foo.valid?
+    end
+    
   end
 
   # TODO: figure out how to test ActiveRecord reload! without a database