Helpers to get target VMs in commands from parameters

lib/vagrant.rb

@@ -12,6 +12,7 @@ module Vagrant
   module Command
     autoload :Base,      'vagrant/command/base'
     autoload :GroupBase, 'vagrant/command/group_base'
+    autoload :Helpers,   'vagrant/command/helpers'
   end
 
   # The source root is the path to the root directory of

lib/vagrant/command/base.rb

@@ -18,6 +18,7 @@ module Vagrant
     # view the various Vagrant commands, which are relatively simple.
     class Base < Thor::Group
       include Thor::Actions
+      include Helpers
 
       attr_reader :env
 

lib/vagrant/command/group_base.rb

@@ -10,6 +10,7 @@ module Vagrant
     # then use {Base} instead.
     class GroupBase < Thor
       include Thor::Actions
+      include Helpers
 
       attr_reader :env
 

lib/vagrant/command/helpers.rb

@@ -0,0 +1,20 @@
+module Vagrant
+  module Command
+    module Helpers
+      # This returns an array of {VM} objects depending on the arguments
+      # given to the command.
+      def target_vms
+        if env.multivm?
+          return env.vms if !self.name
+          vm = env.vms[self.name.to_sym]
+          raise VMNotFoundError.new("A VM by the name of `#{self.name}` was not found.") if !vm
+        else
+          raise MultiVMEnvironmentRequired.new("A multi-vm environment is required for name specification to a command.") if self.name
+          vm = env.vms.values.first
+        end
+
+        [vm]
+      end
+    end
+  end
+end

lib/vagrant/command/status.rb

@@ -10,16 +10,9 @@ module Vagrant
       end
 
       def route
-        if env.multivm?
+        vms = target_vms
-          return show_mulitvm if !self.name
+        show_multivm if vms.length > 1
-          vm = env.vms[self.name.to_sym]
+        show_single(vms.first)
-          raise VMNotFoundError.new("A VM by the name of `#{self.name}` was not found.") if !vm
-        else
-          raise MultiVMEnvironmentRequired.new("A multi-vm environment is required for name specification to a command.") if self.name
-          vm = env.vms.values.first
-        end
-
-        show_single(vm)
       end
 
       protected

test/vagrant/command/helpers_test.rb

@@ -0,0 +1,63 @@
+require "test_helper"
+
+class CommandHelpersTest < Test::Unit::TestCase
+  setup do
+    @module = Vagrant::Command::Helpers
+    @command = Class.new(Vagrant::Command::Base) do
+      argument :name, :optional => true, :type => :string
+    end
+  end
+
+  def command(args, env)
+    @command.new(args, {}, { :env => env })
+  end
+
+  context "vms from args" do
+    setup do
+      @env = mock_environment
+    end
+
+    context "without multivm" do
+      setup do
+        @env.stubs(:vms).returns({ :one => 1 })
+      end
+
+      should "raise an exception if a name is specified" do
+        instance = command(["foo"], @env)
+        assert_raises(Vagrant::MultiVMEnvironmentRequired) {
+          instance.target_vms
+        }
+      end
+
+      should "return the VM if no name is specified" do
+        instance = command([], @env)
+        assert_nothing_raised {
+          assert_equal @env.vms.values, instance.target_vms
+        }
+      end
+    end
+
+    context "with multivm" do
+      setup do
+        @env.stubs(:vms).returns(:one => 1, :two => 2)
+      end
+
+      should "return all the VMs if no name is specified" do
+        instance = command([], @env)
+        assert_equal @env.vms, instance.target_vms
+      end
+
+      should "return only the specified VM if a name is given" do
+        instance = command(["one"], @env)
+        assert_equal @env.vms[:one], instance.target_vms.first
+      end
+
+      should "raise an exception if an invalid name is given" do
+        instance = command(["foo"], @env)
+        assert_raises(Vagrant::VMNotFoundError) {
+          instance.target_vms
+        }
+      end
+    end
+  end
+end