Add simple support for ActiveModel's StateMachine for ActiveRecord

activemodel/lib/active_model/state_machine.rb

@@ -5,12 +5,9 @@ module StateMachine
     autoload :State, 'active_model/state_machine/state'
     autoload :StateTransition, 'active_model/state_machine/state_transition'
 
-    class InvalidTransition < Exception
-    end
+    extend ActiveSupport::Concern
 
-    def self.included(base)
-      require 'active_model/state_machine/machine'
-      base.extend ClassMethods
+    class InvalidTransition < Exception
     end
 
     module ClassMethods

activemodel/lib/active_model/state_machine/event.rb

-require 'active_model/state_machine/state_transition'
-
 module ActiveModel
   module StateMachine
     class Event
@@ -53,12 +51,12 @@ def update(options = {}, &block)
         self
       end
 
-    private
-      def transitions(trans_opts)
-        Array(trans_opts[:from]).each do |s|
-          @transitions << StateTransition.new(trans_opts.merge({:from => s.to_sym}))
+      private
+        def transitions(trans_opts)
+          Array(trans_opts[:from]).each do |s|
+            @transitions << StateTransition.new(trans_opts.merge({:from => s.to_sym}))
+          end
         end
-      end
     end
   end
 end

activemodel/lib/active_model/state_machine/machine.rb

-require 'active_model/state_machine/state'
-require 'active_model/state_machine/event'
-
 module ActiveModel
   module StateMachine
     class Machine
@@ -57,22 +54,22 @@ def current_state_variable
         "@#{@name}_current_state"
       end
 
-    private
-      def state(name, options = {})
-        @states << (state_index[name] ||= State.new(name, :machine => self)).update(options)
-      end
+      private
+        def state(name, options = {})
+          @states << (state_index[name] ||= State.new(name, :machine => self)).update(options)
+        end
 
-      def event(name, options = {}, &block)
-        (@events[name] ||= Event.new(self, name)).update(options, &block)
-      end
+        def event(name, options = {}, &block)
+          (@events[name] ||= Event.new(self, name)).update(options, &block)
+        end
 
-      def event_fired_callback
-        @event_fired_callback ||= (@name == :default ? '' : "#{@name}_") + 'event_fired'
-      end
+        def event_fired_callback
+          @event_fired_callback ||= (@name == :default ? '' : "#{@name}_") + 'event_fired'
+        end
 
-      def event_failed_callback
-        @event_failed_callback ||= (@name == :default ? '' : "#{@name}_") + 'event_failed'
-      end
+        def event_failed_callback
+          @event_failed_callback ||= (@name == :default ? '' : "#{@name}_") + 'event_failed'
+        end
     end
   end
 end

activemodel/lib/active_model/state_machine/state_transition.rb

@@ -18,7 +18,7 @@ def perform(obj)
           true
         end
       end
-      
+
       def execute(obj, *args)
         case @on_transition
         when Symbol, String

activerecord/lib/active_record.rb

@@ -65,6 +65,7 @@ def self.load_all!
   autoload :SchemaDumper, 'active_record/schema_dumper'
   autoload :Serialization, 'active_record/serialization'
   autoload :SessionStore, 'active_record/session_store'
+  autoload :StateMachine, 'active_record/state_machine'
   autoload :TestCase, 'active_record/test_case'
   autoload :Timestamp, 'active_record/timestamp'
   autoload :Transactions, 'active_record/transactions'

activerecord/lib/active_record/state_machine.rb

+module ActiveRecord
+  module StateMachine #:nodoc:
+    extend ActiveSupport::Concern
+    include ActiveModel::StateMachine
+
+    included do
+      before_validation :set_initial_state
+      validates_presence_of :state
+    end
+
+    protected
+      def write_state(state_machine, state)
+        update_attributes! :state => state.to_s
+      end
+
+      def read_state(state_machine)
+        self.state.to_sym
+      end
+
+      def set_initial_state
+        self.state ||= self.class.state_machine.initial_state.to_s
+      end
+  end
+end

activerecord/test/cases/state_machine_test.rb

+require 'cases/helper'
+require 'models/traffic_light'
+
+class StateMachineTest < ActiveRecord::TestCase
+  def setup
+    @light = TrafficLight.create!
+  end
+
+  test "states initial state" do
+    assert @light.off?
+    assert_equal :off, @light.current_state
+  end
+
+  test "transition to a valid state" do
+    @light.reset
+    assert @light.red?
+    assert_equal :red, @light.current_state
+
+    @light.green_on
+    assert @light.green?
+    assert_equal :green, @light.current_state
+  end
+
+  test "transition does not persist state" do
+    @light.reset
+    assert_equal :red, @light.current_state
+    @light.reload
+    assert_equal "off", @light.state
+  end
+
+  test "transition does persists state" do
+    @light.reset!
+    assert_equal :red, @light.current_state
+    @light.reload
+    assert_equal "red", @light.state
+  end
+
+  test "transition to an invalid state" do
+    assert_raise(ActiveModel::StateMachine::InvalidTransition) { @light.yellow_on }
+    assert_equal :off, @light.current_state
+  end
+end

activerecord/test/models/traffic_light.rb

+class TrafficLight < ActiveRecord::Base
+  include ActiveRecord::StateMachine
+
+  state_machine do
+    state :off
+
+    state :red
+    state :green
+    state :yellow
+
+    event :red_on do
+      transitions :to => :red, :from => [:yellow]
+    end
+
+    event :green_on do
+      transitions :to => :green, :from => [:red]
+    end
+
+    event :yellow_on do
+      transitions :to => :yellow, :from => [:green]
+    end
+
+    event :reset do
+      transitions :to => :red, :from => [:off]
+    end
+  end
+end

activerecord/test/schema/schema.rb

@@ -448,6 +448,13 @@ def create_table(*args, &block)
     t.integer :pet_id, :integer
   end
 
+  create_table :traffic_lights, :force => true do |t|
+    t.string   :location
+    t.string   :state
+    t.datetime :created_at
+    t.datetime :updated_at
+  end
+
   create_table :treasures, :force => true do |t|
     t.column :name, :string
     t.column :looter_id, :integer