From 80c05460ab20249ee3b6c44a0960ef48dc411fb2 Mon Sep 17 00:00:00 2001
From: Chris Roberts <croberts@hashicorp.com>
Date: Mon, 16 Dec 2019 17:10:12 -0800
Subject: [PATCH] Call hooks before and after each action if they are available

---
 lib/vagrant/action/warden.rb            |   9 ++
 test/unit/vagrant/action/warden_test.rb | 120 +++++++++++++++---------
 2 files changed, 85 insertions(+), 44 deletions(-)

diff --git a/lib/vagrant/action/warden.rb b/lib/vagrant/action/warden.rb
index 222d19ea3..105b6f928 100644
--- a/lib/vagrant/action/warden.rb
+++ b/lib/vagrant/action/warden.rb
@@ -47,7 +47,16 @@ module Vagrant
           raise Errors::VagrantInterrupt if env[:interrupted]
           action = @actions.shift
           @logger.info("Calling IN action: #{action}")
+
+          if !action.is_a?(Proc) && env[:hook]
+            hook_name = action.class.name.split("::").last.
+              gsub(/([a-z])([A-Z])/, '\1_\2').gsub('-', '_').downcase
+          end
+
+          env[:hook].call("before_#{hook_name}".to_sym) if hook_name
           @stack.unshift(action).first.call(env)
+          env[:hook].call("after_#{hook_name}".to_sym) if hook_name
+
           raise Errors::VagrantInterrupt if env[:interrupted]
           @logger.info("Calling OUT action: #{action}")
         rescue SystemExit
diff --git a/test/unit/vagrant/action/warden_test.rb b/test/unit/vagrant/action/warden_test.rb
index 9c668e371..d119f871a 100644
--- a/test/unit/vagrant/action/warden_test.rb
+++ b/test/unit/vagrant/action/warden_test.rb
@@ -1,6 +1,48 @@
 require File.expand_path("../../../base", __FILE__)
 
 describe Vagrant::Action::Warden do
+  class ActionOne
+    def initialize(app, env)
+      @app = app
+    end
+
+    def call(env)
+      @app.call(env)
+    end
+
+    def recover(env)
+      env[:recover] << 1
+    end
+  end
+
+  class ActionTwo
+    def initialize(app, env)
+      @app = app
+    end
+
+    def call(env)
+      @app.call(env)
+    end
+
+    def recover(env)
+      env[:recover] << 2
+    end
+  end
+
+  class ExitAction
+    def initialize(app, env)
+      @app = app
+    end
+
+    def call(env)
+      @app.call(env)
+    end
+
+    def recover(env)
+      env[:recover] = true
+    end
+  end
+
   let(:data) { { data: [] } }
   let(:instance) { described_class.new }
 
@@ -18,38 +60,10 @@ describe Vagrant::Action::Warden do
   end
 
   it "starts a recovery sequence when an exception is raised" do
-    class Action
-      def initialize(app, env)
-        @app = app
-      end
-
-      def call(env)
-        @app.call(env)
-      end
-
-      def recover(env)
-        env[:recover] << 1
-      end
-    end
-
-    class ActionTwo
-      def initialize(app, env)
-        @app = app
-      end
-
-      def call(env)
-        @app.call(env)
-      end
-
-      def recover(env)
-        env[:recover] << 2
-      end
-    end
-
     error_proc = Proc.new { raise "ERROR!" }
 
     data     = { recover: [] }
-    instance = described_class.new([Action, ActionTwo, error_proc], data)
+    instance = described_class.new([ActionOne, ActionTwo, error_proc], data)
 
     # The error should be raised back up
     expect { instance.call(data) }.
@@ -63,25 +77,11 @@ describe Vagrant::Action::Warden do
   end
 
   it "does not do a recovery sequence if SystemExit is raised" do
-    class Action
-      def initialize(app, env)
-        @app = app
-      end
-
-      def call(env)
-        @app.call(env)
-      end
-
-      def recover(env)
-        env[:recover] = true
-      end
-    end
-
     # Make a proc that just calls "abort" which raises a
     # SystemExit exception.
     error_proc = Proc.new { abort }
 
-    instance = described_class.new([Action, error_proc], data)
+    instance = described_class.new([ExitAction, error_proc], data)
 
     # The SystemExit should come through
     expect { instance.call(data) }.to raise_error(SystemExit)
@@ -89,4 +89,36 @@ describe Vagrant::Action::Warden do
     # The recover should not have been called
     expect(data.key?(:recover)).not_to be
   end
+
+  context "when hook is defined" do
+    let(:hook) { double("hook") }
+
+    before do
+      data[:hook] = hook
+      allow(hook).to receive(:call)
+    end
+
+    it "should receive a before hook call" do
+      expect(hook).to receive(:call).with(:before_action_one)
+      described_class.new([ActionOne], data).call(data)
+    end
+
+    it "should receive an after hook call" do
+      expect(hook).to receive(:call).with(:after_action_one)
+      described_class.new([ActionOne], data).call(data)
+    end
+
+    it "should not receive any hook calls for proc instances" do
+      expect(hook).not_to receive(:call)
+      described_class.new([proc{|*_| :testing }], data).call(data)
+    end
+
+    it "should receive before and after calls for each class" do
+      expect(hook).to receive(:call).with(:before_action_one)
+      expect(hook).to receive(:call).with(:after_action_one)
+      expect(hook).to receive(:call).with(:before_action_two)
+      expect(hook).to receive(:call).with(:after_action_two)
+      described_class.new([ActionOne, proc{|*_| :testing }, ActionTwo], data).call(data)
+    end
+  end
 end