From 46878f9c3e73604d437c41f7768b1933b1cbc627 Mon Sep 17 00:00:00 2001
From: Adrien Siami <adrien.siami@dimelo.com>
Date: Tue, 3 Sep 2013 13:13:53 +0200
Subject: [PATCH] Allow a block to be called when running a transition

---
 lib/aasm/aasm.rb                              | 10 +++++---
 lib/aasm/base.rb                              |  8 +++----
 .../persistence/active_record_persistence.rb  |  2 +-
 spec/unit/transition_spec.rb                  | 23 +++++++++++++++++++
 4 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/lib/aasm/aasm.rb b/lib/aasm/aasm.rb
index 42c9b26..09385ca 100644
--- a/lib/aasm/aasm.rb
+++ b/lib/aasm/aasm.rb
@@ -141,7 +141,7 @@ module AASM
 
 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)
@@ -151,7 +151,7 @@ private
       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
@@ -174,9 +174,13 @@ private
     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
diff --git a/lib/aasm/base.rb b/lib/aasm/base.rb
index d0f01d4..1226549 100644
--- a/lib/aasm/base.rb
+++ b/lib/aasm/base.rb
@@ -57,12 +57,12 @@ module AASM
         aasm.may_fire_event?(name, *args)
       end
 
-      @clazz.send(:define_method, "#{name.to_s}!") do |*args|
-        aasm_fire_event(name, {:persist => true}, *args)
+      @clazz.send(:define_method, "#{name.to_s}!") do |*args, &block|
+        aasm_fire_event(name, {:persist => true}, *args, &block)
       end
 
-      @clazz.send(:define_method, "#{name.to_s}") do |*args|
-        aasm_fire_event(name, {:persist => false}, *args)
+      @clazz.send(:define_method, "#{name.to_s}") do |*args, &block|
+        aasm_fire_event(name, {:persist => false}, *args, &block)
       end
     end
 
diff --git a/lib/aasm/persistence/active_record_persistence.rb b/lib/aasm/persistence/active_record_persistence.rb
index 09ae4fd..0a4dbab 100644
--- a/lib/aasm/persistence/active_record_persistence.rb
+++ b/lib/aasm/persistence/active_record_persistence.rb
@@ -131,7 +131,7 @@ module AASM
           aasm.enter_initial_state if send(self.class.aasm_column).blank?
         end
 
-        def aasm_fire_event(name, options, *args)
+        def aasm_fire_event(name, options, *args, &block)
           self.class.transaction(:requires_new => true) do
             super
           end
diff --git a/spec/unit/transition_spec.rb b/spec/unit/transition_spec.rb
index 2e03344..ca60484 100644
--- a/spec/unit/transition_spec.rb
+++ b/spec/unit/transition_spec.rb
@@ -26,6 +26,29 @@ describe 'transitions' do
     silencer.should be_smiling
   end
 
+  it 'should call the block when success' do
+    silencer = Silencer.new
+    success = false
+    lambda {
+      silencer.smile_any! do
+        success = true
+      end
+    }.should change { success }.to(true)
+  end
+
+  it 'should not call the block when failure' do
+    silencer = Silencer.new
+    success = false
+    lambda {
+      silencer.smile! do
+        success = true
+      end
+    }.should_not change { success }.to(true)
+  end
+
+end
+
+describe 'blocks' do
 end
 
 describe AASM::Transition do