Implemented "around" callbacks on VM for convenient `before_foo` and `after_foo` callbacks

lib/vagrant/actions/import.rb

@@ -2,14 +2,12 @@ module Vagrant
   module Actions
     class Import < Base
       def execute!
-        @vm.invoke_callback(:before_import)
-
-        Busy.busy do
-          logger.info "Importing base VM (#{Vagrant.config[:vm][:base]})..."
-          @vm.vm = VirtualBox::VM.import(File.expand_path(Vagrant.config[:vm][:base]))
+        @vm.invoke_around_callback(:import) do
+          Busy.busy do
+            logger.info "Importing base VM (#{Vagrant.config[:vm][:base]})..."
+            @vm.vm = VirtualBox::VM.import(File.expand_path(Vagrant.config[:vm][:base]))
+          end
         end
-
-        @vm.invoke_callback(:after_import)
       end
     end
   end

lib/vagrant/actions/start.rb

@@ -2,20 +2,18 @@ module Vagrant
   module Actions
     class Start < Base
       def execute!
-        @vm.invoke_callback(:before_boot)
+        @vm.invoke_around_callback(:boot) do
+          # Startup the VM
+          boot
 
-        # Startup the VM
-        boot
-
-        # Wait for it to complete booting, or error if we could
-        # never detect it booted up successfully
-        if !wait_for_boot
-          error_and_exit(<<-error)
+          # Wait for it to complete booting, or error if we could
+          # never detect it booted up successfully
+          if !wait_for_boot
+            error_and_exit(<<-error)
 Failed to connect to VM! Failed to boot?
 error
+          end
         end
-
-        @vm.invoke_callback(:after_boot)
       end
 
       def boot

lib/vagrant/vm.rb

@@ -98,6 +98,15 @@ error
       end
     end
 
+    # Invokes an "around callback" which invokes before_name and
+    # after_name for the given callback name, yielding a block between
+    # callback invokations.
+    def invoke_around_callback(name, *args)
+      invoke_callback("before_#{name}".to_sym, *args)
+      yield
+      invoke_callback("after_#{name}".to_sym, *args)
+    end
+
     def invoke_callback(name, *args)
       # Attempt to call the method for the callback on each of the
       # actions

test/test_helper.rb

@@ -66,6 +66,7 @@ class Test::Unit::TestCase
     @mock_vm.stubs(:vm).returns(@vm)
     @mock_vm.stubs(:vm=)
     @mock_vm.stubs(:invoke_callback)
+    @mock_vm.stubs(:invoke_around_callback).yields
     @action = action_klass.new(@mock_vm)
 
     [@mock_vm, @vm, @action]

test/vagrant/actions/import_test.rb

@@ -12,11 +12,8 @@ class ImportActionTest < Test::Unit::TestCase
     @import.execute!
   end
 
-  should "invoke before/after callbacks around the import" do
-    callback_seq = sequence("callback_seq")
-    @mock_vm.expects(:invoke_callback).with(:before_import).once.in_sequence(callback_seq)
-    VirtualBox::VM.expects(:import).once.in_sequence(callback_seq)
-    @mock_vm.expects(:invoke_callback).with(:after_import).once.in_sequence(callback_seq)
+  should "invoke an around callback around the import" do
+    @mock_vm.expects(:invoke_around_callback).with(:import).once
     @import.execute!
   end
 

test/vagrant/actions/start_test.rb

@@ -8,12 +8,11 @@ class StartActionTest < Test::Unit::TestCase
   end
 
   context "execution" do
-    should "invoke before callback, boot, and invoke the after callback" do
+    should "invoke the 'boot' around callback" do
       boot_seq = sequence("boot_seq")
-      @mock_vm.expects(:invoke_callback).with(:before_boot).once.in_sequence(boot_seq)
+      @mock_vm.expects(:invoke_around_callback).with(:boot).once.in_sequence(boot_seq).yields
       @action.expects(:boot).in_sequence(boot_seq)
       @action.expects(:wait_for_boot).returns(true).in_sequence(boot_seq)
-      @mock_vm.expects(:invoke_callback).with(:after_boot).once.in_sequence(boot_seq)
       @action.execute!
     end
 
@@ -21,7 +20,6 @@ class StartActionTest < Test::Unit::TestCase
       fail_boot_seq = sequence("fail_boot_seq")
       @action.expects(:boot).once.in_sequence(fail_boot_seq)
       @action.expects(:wait_for_boot).returns(false).in_sequence(fail_boot_seq)
-      @action.expects(:invoke_callback).with(:after_boot).never
       @action.expects(:error_and_exit).once.in_sequence(fail_boot_seq)
       @action.execute!
     end

test/vagrant/vm_test.rb

@@ -16,6 +16,26 @@ class VMTest < Test::Unit::TestCase
       @vm = Vagrant::VM.new(@mock_vm)
     end
 
+    context "around callbacks" do
+      should "invoke before/after_name for around callbacks" do
+        block_obj = mock("block_obj")
+        around_seq = sequence("around_seq")
+        @vm.expects(:invoke_callback).with(:before_foo).once.in_sequence(around_seq)
+        block_obj.expects(:foo).once.in_sequence(around_seq)
+        @vm.expects(:invoke_callback).with(:after_foo).once.in_sequence(around_seq)
+
+        @vm.invoke_around_callback(:foo) do
+          block_obj.foo
+        end
+      end
+
+      should "forward arguments to invoke_callback" do
+        @vm.expects(:invoke_callback).with(:before_foo, "foo").once
+        @vm.expects(:invoke_callback).with(:after_foo, "foo").once
+        @vm.invoke_around_callback(:foo, "foo") do; end
+      end
+    end
+
     should "not invoke callback on actions which don't respond to it" do
       action = mock("action")
       action.stubs(:respond_to?).with(:foo).returns(false)