From 7691b28c5377424d68392dd6a18fdab8d6a30ad3 Mon Sep 17 00:00:00 2001
From: Mitchell Hashimoto <mitchell.hashimoto@gmail.com>
Date: Tue, 15 Jun 2010 22:31:51 -0700
Subject: [PATCH] `vagrant ssh` now takes an `--execute` (or `-e`) param for
 executing commands [closes GH-95]

---
 lib/vagrant/commands/ssh.rb       | 38 +++++++++++++++++++--
 test/vagrant/commands/ssh_test.rb | 57 +++++++++++++++++++++++++++++++
 2 files changed, 93 insertions(+), 2 deletions(-)

diff --git a/lib/vagrant/commands/ssh.rb b/lib/vagrant/commands/ssh.rb
index 3a8c6d243..2522ba2be 100644
--- a/lib/vagrant/commands/ssh.rb
+++ b/lib/vagrant/commands/ssh.rb
@@ -11,7 +11,34 @@ module Vagrant
 
       def execute(args=[])
         args = parse_options(args)
-        ssh_connect(args[0])
+        if !options[:execute].empty?
+          vms = args[0] ? {args[0] => env.vms[args[0].to_sym]} : env.vms
+          vms.each do |name, vm|
+            ssh_execute(name, vm)
+          end
+        else
+          ssh_connect(args[0])
+        end
+      end
+
+      def ssh_execute(name, vm)
+        if vm.nil?
+          error_and_exit(:unknown_vm, :vm => name)
+          return # for tests
+        elsif !vm.created?
+          error_and_exit(:environment_not_created)
+          return
+        end
+
+        vm.ssh.execute do |ssh|
+          options[:execute].each do |command|
+            vm.env.logger.info("Execute: #{command}")
+            ssh.exec!(command) do |channel, type, data|
+              # TODO: Exit status checking?
+              vm.env.logger.info("#{type}: #{data}")
+            end
+          end
+        end
       end
 
       def ssh_connect(name)
@@ -37,7 +64,14 @@ module Vagrant
       end
 
       def options_spec(opts)
-        opts.banner = "Usage: vagrant ssh"
+        opts.banner = "Usage: vagrant ssh [--execute COMMAND]"
+
+        # Defaults
+        options[:execute] = []
+
+        opts.on("-e", "--execute COMMAND", "A command to execute. Multiple -e's may be specified.") do |value|
+          options[:execute] << value
+        end
       end
     end
   end
diff --git a/test/vagrant/commands/ssh_test.rb b/test/vagrant/commands/ssh_test.rb
index 84d80e255..437325399 100644
--- a/test/vagrant/commands/ssh_test.rb
+++ b/test/vagrant/commands/ssh_test.rb
@@ -18,6 +18,63 @@ class CommandsSSHTest < Test::Unit::TestCase
       @instance.expects(:ssh_connect).with(nil).once
       @instance.execute
     end
+
+    should "execute if commands are given" do
+      @env.stubs(:vms).returns(:foo => mock("foo"))
+      @instance.expects(:ssh_execute).with("foo", @env.vms[:foo]).once
+      @instance.execute(["foo","-e","bar"])
+    end
+
+    should "execute over every VM if none is specified with a command" do
+      vms = {
+        :foo => mock("foo"),
+        :bar => mock("bar")
+      }
+
+      @env.stubs(:vms).returns(vms)
+      vms.each do |key, vm|
+        @instance.expects(:ssh_execute).with(key, vm).once
+      end
+
+      @instance.execute(["--execute", "bar"])
+    end
+  end
+
+  context "ssh executing" do
+    setup do
+      @name = :foo
+
+      @ssh_conn = mock("connection")
+      @ssh_conn.stubs(:exec!)
+
+      @ssh = mock("ssh")
+      @ssh.stubs(:execute).yields(@ssh_conn)
+
+      @vm = mock("vm")
+      @vm.stubs(:created?).returns(true)
+      @vm.stubs(:ssh).returns(@ssh)
+    end
+
+    should "error and exit if invalid VM given" do
+      @instance.expects(:error_and_exit).with(:unknown_vm, :vm => @name).once
+      @instance.ssh_execute(@name, nil)
+    end
+
+    should "error and exit if VM isn't created" do
+      @vm.stubs(:created?).returns(false)
+      @instance.expects(:error_and_exit).with(:environment_not_created).once
+      @instance.ssh_execute(@name, @vm)
+    end
+
+    should "execute each command" do
+      commands = [:a,:b,:c]
+      @instance.stubs(:options).returns(:execute => commands)
+      commands.each do |cmd|
+        @ssh_conn.expects(:exec!).with(cmd).once
+      end
+
+      @instance.ssh_execute(@name, @vm)
+    end
   end
 
   context "ssh connecting" do