Allow subclassing AASM core classes

lib/aasm.rb

@@ -3,6 +3,7 @@
 require 'aasm/configuration'
 require 'aasm/base'
 require 'aasm/dsl_helper'
+require 'aasm/transition_builder'
 require 'aasm/instance_base'
 require 'aasm/core/transition'
 require 'aasm/core/event'

lib/aasm/base.rb

@@ -10,6 +10,7 @@ def initialize(klass, name, state_machine, options={}, &block)
       @name = name
       # @state_machine = klass.aasm(@name).state_machine
       @state_machine = state_machine
+      @state_machine.implementation = self
       @state_machine.config.column ||= (options[:column] || default_column).to_sym
       # @state_machine.config.column = options[:column].to_sym if options[:column] # master
       @options = options
@@ -196,6 +197,18 @@ def from_states_for_state(state, options={})
       end
     end
 
+    def aasm_state_class
+      AASM::Core::State
+    end
+
+    def aasm_event_class
+      AASM::Core::Event
+    end
+
+    def aasm_transition_class
+      AASM::Core::Transition
+    end
+
     private
 
     def default_column

lib/aasm/core/event.rb

@@ -3,6 +3,7 @@
 module AASM::Core
   class Event
     include AASM::DslHelper
+    include AASM::TransitionBuilder
 
     attr_reader :name, :state_machine, :options, :default_display_name
 
@@ -17,17 +18,7 @@ def initialize(name, state_machine, options = {}, &block)
 
       # from aasm4
       @options = options # QUESTION: .dup ?
-      add_options_from_dsl(@options, [
-        :after,
-        :after_commit,
-        :after_transaction,
-        :before,
-        :before_transaction,
-        :ensure,
-        :error,
-        :before_success,
-        :success,
-      ], &block) if block
+      add_options_from_dsl(@options, dsl_option_keys, &block) if block
     end
 
     # called internally by Ruby 1.9 after clone()
@@ -96,12 +87,12 @@ def ==(event)
     def transitions(definitions=nil, &block)
       if definitions # define new transitions
         # Create a separate transition for each from-state to the given state
-        Array(definitions[:from]).each do |s|
-          @transitions << AASM::Core::Transition.new(self, attach_event_guards(definitions.merge(:from => s.to_sym)), &block)
+        Array(definitions[:from]).each do |from|
+          build_transition(definitions, from, &block)
         end
         # Create a transition if :to is specified without :from (transitions from ANY state)
         if !definitions[:from] && definitions[:to]
-          @transitions << AASM::Core::Transition.new(self, attach_event_guards(definitions), &block)
+          build_transition(definitions, &block)
         end
       end
       @transitions
@@ -117,18 +108,6 @@ def to_s
 
   private
 
-    def attach_event_guards(definitions)
-      unless @guards.empty?
-        given_guards = Array(definitions.delete(:guard) || definitions.delete(:guards) || definitions.delete(:if))
-        definitions[:guards] = @guards + given_guards # from aasm4
-      end
-      unless @unless.empty?
-        given_unless = Array(definitions.delete(:unless))
-        definitions[:unless] = given_unless + @unless
-      end
-      definitions
-    end
-
     def _fire(obj, options={}, to_state=::AASM::NO_VALUE, *args)
       result = options[:test_only] ? false : nil
       clear_failed_callbacks
@@ -142,7 +121,7 @@ def _fire(obj, options={}, to_state=::AASM::NO_VALUE, *args)
         args.unshift(to_state)
         to_state = nil
       end
-      
+
       # nop, to_state is a valid to-state
 
       transitions.each do |transition|
@@ -174,5 +153,19 @@ def invoke_callbacks(code, record, args)
              .with_default_return_value(false)
              .invoke
     end
+
+    def dsl_option_keys
+      [
+        :after,
+        :after_commit,
+        :after_transaction,
+        :before,
+        :before_transaction,
+        :ensure,
+        :error,
+        :before_success,
+        :success,
+      ]
+    end
   end
 end # AASM

lib/aasm/core/transition.rb

@@ -8,7 +8,7 @@ class Transition
     alias_method :options, :opts
 
     def initialize(event, opts, &block)
-      add_options_from_dsl(opts, [:on_transition, :guard, :after, :success], &block) if block
+      add_options_from_dsl(opts, dsl_option_keys, &block) if block
 
       @event = event
       @from = opts[:from]
@@ -79,5 +79,8 @@ def _fire_callbacks(code, record, args)
       Invoker.new(code, record, args).invoke
     end
 
+    def dsl_option_keys
+      [:on_transition, :guard, :after, :success]
+    end
   end
 end # AASM

lib/aasm/state_machine.rb

@@ -2,7 +2,7 @@ module AASM
   class StateMachine
     # the following four methods provide the storage of all state machines
 
-    attr_accessor :states, :events, :initial_state, :config, :name, :global_callbacks
+    attr_accessor :states, :events, :initial_state, :config, :name, :implementation, :global_callbacks
 
     def initialize(name)
       @initial_state = nil
@@ -28,11 +28,15 @@ def add_state(state_name, klass, options)
       # allow reloading, extending or redefining a state
       @states.delete(state_name) if @states.include?(state_name)
 
-      @states << AASM::Core::State.new(state_name, klass, self, options)
+      state_class = implementation.aasm_state_class
+      raise ArgumentError, "The class #{state_class} must inherit from AASM::Core::State!" unless state_class.ancestors.include?(AASM::Core::State)
+      @states << state_class.new(state_name, klass, self, options)
     end
 
     def add_event(name, options, &block)
-      @events[name] = AASM::Core::Event.new(name, self, options, &block)
+      event_class = implementation.aasm_event_class
+      raise ArgumentError, "The class #{event_class} must inherit from AASM::Core::Event!" unless event_class.ancestors.include?(AASM::Core::Event)
+      @events[name] = event_class.new(name, self, options, &block)
     end
 
     def add_global_callbacks(name, *callbacks, &block)

lib/aasm/transition_builder.rb

@@ -0,0 +1,28 @@
+module AASM
+  module TransitionBuilder
+
+  private
+
+    def build_transition(definitions, from = nil, &block)
+      transition_class = state_machine.implementation.aasm_transition_class
+      raise ArgumentError, "The class #{transition_class} must inherit from AASM::Core::Transition!" unless transition_class.ancestors.include?(AASM::Core::Transition)
+
+      definitions = definitions.merge(:from => from.to_sym) if from
+
+      @transitions << transition_class.new(self, attach_event_guards(definitions), &block)
+    end
+
+    def attach_event_guards(definitions)
+      unless @guards.empty?
+        given_guards = Array(definitions.delete(:guard) || definitions.delete(:guards) || definitions.delete(:if))
+        definitions[:guards] = @guards + given_guards # from aasm4
+      end
+      unless @unless.empty?
+        given_unless = Array(definitions.delete(:unless))
+        definitions[:unless] = given_unless + @unless
+      end
+      definitions
+    end
+
+  end
+end

spec/models/custom_aasm_base.rb

@@ -0,0 +1,14 @@
+class CustomAasmBase < AASM::Base
+
+  def aasm_state_class
+    CustomState
+  end
+
+  def aasm_event_class
+    CustomEvent
+  end
+
+  def aasm_transition_class
+    CustomTransition
+  end
+end

spec/models/custom_event.rb

@@ -0,0 +1,21 @@
+class CustomEvent < AASM::Core::Event
+  attr_reader :custom_method_args
+
+  def custom_event_method!(value)
+    @custom_method_args = value
+  end
+
+  def some_option
+    options[:some_option]
+  end
+
+  def another_option
+    options[:another_option]
+  end
+
+  private
+
+  def dsl_option_keys
+    super + [:some_option, :another_option]
+  end
+end

spec/models/custom_state.rb

@@ -0,0 +1,6 @@
+class CustomState < AASM::Core::State
+
+  def custom_state_method(value)
+    value * value
+  end
+end

spec/models/custom_transition.rb

@@ -0,0 +1,21 @@
+class CustomTransition < AASM::Core::Transition
+  attr_reader :custom_method_args
+
+  def custom_transition_method!(value)
+    @custom_method_args = value
+  end
+
+  def some_option
+    opts[:some_option]
+  end
+
+  def another_option
+    options[:another_option]
+  end
+
+  private
+
+  def dsl_option_keys
+    super + [:some_option, :another_option]
+  end
+end

spec/models/full_example_with_custom_aasm_base.rb

@@ -0,0 +1,18 @@
+class FullExampleWithCustomAasmBase
+  include AASM
+
+  aasm with_klass: CustomAasmBase do
+    state :initialised, :initial => true
+    state :filled_out
+
+    event :fill_out, :some_option => '-- some event value --' do
+      another_option '-- another event value --'
+      custom_event_method!(41)
+
+      transitions :from => :initialised, :to => :filled_out, :some_option => '-- some transition value --' do
+        another_option '-- another transition value --'
+        custom_transition_method! 42
+      end
+    end
+  end
+end

spec/models/real_world_example_with_custom_aasm_base.rb

@@ -0,0 +1,26 @@
+class RealWorldExampleWithCustomAasmBase
+  include AASM
+
+  class RequiredParamsEvent < AASM::Core::Event
+    def required_params!(*keys)
+      options[:before] ||= []
+      options[:before] << ->(**args) do
+        missing = keys - args.keys
+        raise ArgumentError, "Missing required arguments #{missing.inspect}" unless missing == []
+      end
+    end
+  end
+  class RequiredParams < AASM::Base
+    def aasm_event_class; RequiredParamsEvent; end
+  end
+
+  aasm with_klass: RequiredParams do
+    state :initialised, :initial => true
+    state :filled_out
+
+    event :fill_out do
+      required_params! :user, :quantity, :date
+      transitions :from => :initialised, :to => :filled_out
+    end
+  end
+end

spec/models/simple_example_with_custom_aasm_base.rb

@@ -0,0 +1,12 @@
+class SimpleExampleWithCustomAasmBase
+  include AASM
+
+  aasm with_klass: CustomAasmBase do
+    state :initialised, :initial => true
+    state :filled_out
+
+    event :fill_out do
+      transitions :from => :initialised, :to => :filled_out
+    end
+  end
+end

spec/unit/dsl_with_custom_aasm_base_spec.rb

@@ -0,0 +1,26 @@
+require 'spec_helper'
+
+describe "dsl with custom ASM::Base and custom core classes" do
+
+  let(:example) {FullExampleWithCustomAasmBase.new}
+
+  it 'should create the expected state machine' do
+    aasm = example.aasm(:default)
+
+    state = aasm.states.first
+    expect(state).to be_a(CustomState)
+
+    event = aasm.events.first
+    expect(event).to be_a(CustomEvent)
+    expect(event.some_option).to eq('-- some event value --')
+    expect(event.another_option).to eq(['-- another event value --'])
+    expect(event.custom_method_args).to eq(41)
+
+    transition = event.transitions.first
+    expect(transition).to be_a(CustomTransition)
+    expect(transition.some_option).to eq('-- some transition value --')
+    expect(transition.another_option).to eq(['-- another transition value --'])
+    expect(transition.custom_method_args).to eq(42)
+  end
+
+end

spec/unit/event_spec.rb

@@ -1,7 +1,7 @@
 require 'spec_helper'
 
 describe 'adding an event' do
-  let(:state_machine) { AASM::StateMachine.new(:name) }
+  let(:state_machine) { AASM::StateMachine.new(:name).tap { |sm| AASM::Base.new(nil, :name, sm) } }
   let(:event) do
     AASM::Core::Event.new(:close_order, state_machine, {:success => :success_callback}) do
       before :before_callback
@@ -36,7 +36,7 @@
 end
 
 describe 'transition inspection' do
-  let(:state_machine) { AASM::StateMachine.new(:name) }
+  let(:state_machine) { AASM::StateMachine.new(:name).tap { |sm| AASM::Base.new(nil, :name, sm) } }
   let(:event) do
     AASM::Core::Event.new(:run, state_machine) do
       transitions :to => :running, :from => :sleeping
@@ -61,7 +61,7 @@
 end
 
 describe 'transition inspection without from' do
-  let(:state_machine) { AASM::StateMachine.new(:name) }
+  let(:state_machine) { AASM::StateMachine.new(:name).tap { |sm| AASM::Base.new(nil, :name, sm) } }
   let(:event) do
     AASM::Core::Event.new(:run, state_machine) do
       transitions :to => :running
@@ -79,7 +79,7 @@
 end
 
 describe 'firing an event' do
-  let(:state_machine) { AASM::StateMachine.new(:name) }
+  let(:state_machine) { AASM::StateMachine.new(:name).tap { |sm| AASM::Base.new(nil, :name, sm) } }
 
   it 'should return nil if the transitions are empty' do
     obj = double('object', :aasm => double('aasm', :current_state => 'open'))