lib/aasm/aasm.rb in aasm-3.0.23 vs lib/aasm/aasm.rb in aasm-3.0.24

@@ -139,21 +139,21 @@
     aasm.human_state
   end
 
 private
 
-  def aasm_fire_event(event_name, options, *args)
+  def aasm_fire_event(event_name, options, *args, &block)
     event = self.class.aasm_events[event_name]
     begin
       old_state = aasm.state_object_for_name(aasm.current_state)
       old_state.fire_callbacks(:exit, self)
 
       # new event before callback
       event.fire_callbacks(:before, self)
 
       if new_state_name = event.fire(self, *args)
-        fired(event, old_state, new_state_name, options)
+        fired(event, old_state, new_state_name, options, &block)
       else
         failed(event_name, old_state)
       end
     rescue StandardError => e
       event.fire_callbacks(:error, self, e) || raise(e)
@@ -172,12 +172,16 @@
     new_state.fire_callbacks(:enter, self)
 
     persist_successful = true
     if persist
       persist_successful = aasm.set_current_state_with_persistence(new_state_name)
-      event.fire_callbacks(:success, self) if persist_successful
+      if persist_successful
+        yield if block_given?
+        event.fire_callbacks(:success, self)
+      end
     else
       aasm.current_state = new_state_name
+      yield if block_given?
     end
 
     if persist_successful
       old_state.fire_callbacks(:after_exit, self)
       new_state.fire_callbacks(:after_enter, self)