From 687b925d2e9bb41d6c18762b9fe57fbe8051f2e6 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Thu, 27 May 2010 22:54:11 -0700
Subject: [PATCH] Concept of a 'primary' VM in a multi-VM environment added.
 This VM will be the default for actions in a multi-VM environment.

---
 lib/vagrant/config.rb            |  5 ++--
 lib/vagrant/environment.rb       | 17 ++++++++++-
 test/vagrant/config_test.rb      |  7 ++++-
 test/vagrant/environment_test.rb | 49 ++++++++++++++++++++++++++++++++
 4 files changed, 74 insertions(+), 4 deletions(-)

diff --git a/lib/vagrant/config.rb b/lib/vagrant/config.rb
index 016caf120..64c403c4d 100644
--- a/lib/vagrant/config.rb
+++ b/lib/vagrant/config.rb
@@ -155,8 +155,9 @@ module Vagrant
         @defined_vms ||= {}
       end
 
-      def define(name, &block)
-        defined_vms[name.to_sym] = block
+      def define(name, options=nil, &block)
+        options ||= {}
+        defined_vms[name.to_sym] = options.merge({:config_proc => block})
       end
 
       def shift(orig, rsync)
diff --git a/lib/vagrant/environment.rb b/lib/vagrant/environment.rb
index 57d2f5639..3258c9b77 100644
--- a/lib/vagrant/environment.rb
+++ b/lib/vagrant/environment.rb
@@ -101,6 +101,18 @@ module Vagrant
       @vms ||= {}
     end
 
+    # Returns the primray VM associated with this environment
+    def primary_vm
+      return vms.values.first if !multivm?
+      return parent.primary_vm if parent
+
+      config.vm.defined_vms.each do |name, options|
+        return vms[name] if options[:primary]
+      end
+
+      nil
+    end
+
     # Returns a boolean whether this environment represents a multi-VM
     # environment or not. This will work even when called on child
     # environments.
@@ -165,7 +177,10 @@ module Vagrant
 
       # If this environment represents some VM in a multi-VM environment,
       # we push that VM's configuration onto the config_queue.
-      config_queue << parent.config.vm.defined_vms[vm_name] if vm_name
+      if vm_name
+        vm_data = parent.config.vm.defined_vms[vm_name] || {}
+        config_queue << vm_data[:config_proc]
+      end
 
       # Clear out the old data
       Config.reset!(self)
diff --git a/test/vagrant/config_test.rb b/test/vagrant/config_test.rb
index a9f34ea9a..a4ffc6d5d 100644
--- a/test/vagrant/config_test.rb
+++ b/test/vagrant/config_test.rb
@@ -232,7 +232,12 @@ class ConfigTest < Test::Unit::TestCase
 
         proc = Proc.new { foo.call }
         @config.define(:name, &proc)
-        assert_equal proc, @config.defined_vms[:name]
+        assert_equal proc, @config.defined_vms[:name][:config_proc]
+      end
+
+      should "store the options" do
+        @config.define(:name, :set => true)
+        assert @config.defined_vms[:name][:set]
       end
 
       should "not have multi-VMs by default" do
diff --git a/test/vagrant/environment_test.rb b/test/vagrant/environment_test.rb
index 2b633fa4e..fcdbd66b9 100644
--- a/test/vagrant/environment_test.rb
+++ b/test/vagrant/environment_test.rb
@@ -124,6 +124,55 @@ class EnvironmentTest < Test::Unit::TestCase
     end
   end
 
+  context "primary VM helper" do
+    setup do
+      @env = mock_environment
+      @env.stubs(:multivm?).returns(true)
+    end
+
+    should "return the first VM if not multivm" do
+      result = mock("result")
+
+      @env.stubs(:multivm?).returns(false)
+      @env.stubs(:vms).returns({:default => result})
+
+      assert_equal result, @env.primary_vm
+    end
+
+    should "call and return the primary VM from the parent if has one" do
+      result = mock("result")
+      parent = mock("parent")
+      parent.expects(:primary_vm).returns(result)
+
+      @env.stubs(:parent).returns(parent)
+      assert_equal result, @env.primary_vm
+    end
+
+    should "return nil if no VM is marked as primary" do
+      @env.config.vm.define(:foo)
+      @env.config.vm.define(:bar)
+      @env.config.vm.define(:baz)
+
+      assert @env.primary_vm.nil?
+    end
+
+    should "return the primary VM" do
+      @env.config.vm.define(:foo)
+      @env.config.vm.define(:bar, :primary => true)
+      @env.config.vm.define(:baz)
+
+      result = mock("result")
+      vms = {
+        :foo => :foo,
+        :bar => result,
+        :baz => :baz
+      }
+      @env.stubs(:vms).returns(vms)
+
+      assert_equal result, @env.primary_vm
+    end
+  end
+
   context "multivm? helper" do
     setup do
       @env = mock_environment