From 85b8ccc2ec27b8f657160d8df89a89cdd3d19ebd Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Tue, 11 Jan 2011 19:52:55 -0800
Subject: [PATCH] Provisioners now register themselves with their shortcut name

---
 lib/vagrant/action/vm/provision.rb        | 10 +---------
 lib/vagrant/provisioners/base.rb          | 13 +++++++++++++
 lib/vagrant/provisioners/chef_server.rb   |  2 ++
 lib/vagrant/provisioners/chef_solo.rb     |  2 ++
 lib/vagrant/provisioners/puppet.rb        |  2 ++
 lib/vagrant/provisioners/puppet_server.rb |  2 ++
 test/vagrant/provisioners/base_test.rb    | 22 ++++++++++++++++++++++
 7 files changed, 44 insertions(+), 9 deletions(-)

diff --git a/lib/vagrant/action/vm/provision.rb b/lib/vagrant/action/vm/provision.rb
index 1f059ebff..6977634ea 100644
--- a/lib/vagrant/action/vm/provision.rb
+++ b/lib/vagrant/action/vm/provision.rb
@@ -30,15 +30,7 @@ module Vagrant
             @provisioner = provisioner.new(@env)
             raise Errors::ProvisionInvalidClass if !@provisioner.is_a?(Provisioners::Base)
           elsif provisioner.is_a?(Symbol)
-            # We have a few hard coded provisioners for built-ins
-            mapping = {
-              :chef_solo      => Provisioners::ChefSolo,
-              :chef_server    => Provisioners::ChefServer,
-              :puppet         => Provisioners::Puppet,
-              :puppet_server  => Provisioners::PuppetServer
-            }
-
-            provisioner_klass = mapping[provisioner]
+            provisioner_klass = Provisioners::Base.registered[provisioner]
             raise Errors::ProvisionUnknownType, :provisioner => provisioner.to_s if provisioner_klass.nil?
             @provisioner = provisioner_klass.new(@env)
           end
diff --git a/lib/vagrant/provisioners/base.rb b/lib/vagrant/provisioners/base.rb
index 637531230..40641ed86 100644
--- a/lib/vagrant/provisioners/base.rb
+++ b/lib/vagrant/provisioners/base.rb
@@ -11,6 +11,19 @@ module Vagrant
       # {Vagrant::Action::Environment}
       attr_reader :action_env
 
+      # Registers a provisioner with a given shortcut. This allows that provisioner
+      # to be referenced with the shortcut.
+      #
+      # @param [Symbol] shortcut
+      def self.register(shortcut)
+        registered[shortcut] = self
+      end
+
+      # Returns the provisioner associated with the given shortcut.
+      def self.registered
+        @@registered ||= {}
+      end
+
       def initialize(env)
         @action_env = env
       end
diff --git a/lib/vagrant/provisioners/chef_server.rb b/lib/vagrant/provisioners/chef_server.rb
index 1b886af09..3cc3f3d45 100644
--- a/lib/vagrant/provisioners/chef_server.rb
+++ b/lib/vagrant/provisioners/chef_server.rb
@@ -5,6 +5,8 @@ module Vagrant
     # This class implements provisioning via chef-client, allowing provisioning
     # with a chef server.
     class ChefServer < Chef
+      register :chef_server
+
       def prepare
         raise ChefError, :server_validation_key_required if env.config.chef.validation_key_path.nil?
         raise ChefError, :server_validation_key_doesnt_exist if !File.file?(validation_key_path)
diff --git a/lib/vagrant/provisioners/chef_solo.rb b/lib/vagrant/provisioners/chef_solo.rb
index d152e9c4e..d8fc67167 100644
--- a/lib/vagrant/provisioners/chef_solo.rb
+++ b/lib/vagrant/provisioners/chef_solo.rb
@@ -2,6 +2,8 @@ module Vagrant
   module Provisioners
     # This class implements provisioning via chef-solo.
     class ChefSolo < Chef
+      register :chef_solo
+
       def prepare
         share_cookbook_folders
         share_role_folders
diff --git a/lib/vagrant/provisioners/puppet.rb b/lib/vagrant/provisioners/puppet.rb
index 9f4739abe..2c01e4e24 100644
--- a/lib/vagrant/provisioners/puppet.rb
+++ b/lib/vagrant/provisioners/puppet.rb
@@ -21,6 +21,8 @@ module Vagrant
     end
 
     class Puppet < Base
+      register :puppet
+
       def prepare
         check_manifest_dir
         share_manifests
diff --git a/lib/vagrant/provisioners/puppet_server.rb b/lib/vagrant/provisioners/puppet_server.rb
index 3eccfd149..89decdd0c 100644
--- a/lib/vagrant/provisioners/puppet_server.rb
+++ b/lib/vagrant/provisioners/puppet_server.rb
@@ -19,6 +19,8 @@ module Vagrant
     end
 
     class PuppetServer < Base
+      register :puppet_server
+
       def provision!
         verify_binary("puppetd")
         run_puppetd_client
diff --git a/test/vagrant/provisioners/base_test.rb b/test/vagrant/provisioners/base_test.rb
index 820cdd48b..65fff9c93 100644
--- a/test/vagrant/provisioners/base_test.rb
+++ b/test/vagrant/provisioners/base_test.rb
@@ -1,10 +1,32 @@
 require "test_helper"
 
 class BaseProvisionerTest < Test::Unit::TestCase
+  setup do
+    @klass = Vagrant::Provisioners::Base
+  end
+
   should "include the util class so subclasses have access to it" do
     assert Vagrant::Provisioners::Base.include?(Vagrant::Util)
   end
 
+  context "registering provisioners" do
+    teardown do
+      @klass.registered.delete(:zomg)
+    end
+
+    should "not have unregistered provisioners" do
+      assert_nil @klass.registered[:foo]
+    end
+
+    should "be able to register a provisioner" do
+      foo = Class.new(@klass) do
+        register :zomg
+      end
+
+      assert_equal foo, @klass.registered[:zomg]
+    end
+  end
+
   context "base instance" do
     setup do
       @env = Vagrant::Action::Environment.new(vagrant_env)