Merge pull request #214 from pearkes/add_ssh_wait_command

Add SSH -w command

lib/tugboat/cli.rb

@@ -127,6 +127,10 @@ def images
                   :type => :string,
                   :aliases => [ "-c", "-y" ,],
                   :desc => "Command to run on the droplet"
+    method_option "wait",
+                  :type => :boolean,
+                  :aliases => '-w',
+                  :desc => 'Wait for droplet to become active before trying to SSH'
     def ssh(name=nil)
       Middleware.sequence_ssh_droplet.call({
         "tugboat_action" => __method__,
@@ -138,7 +142,8 @@ def ssh(name=nil)
         "user_droplet_use_private_ip" => options[:use_private_ip],
         "user_droplet_ssh_opts" => options[:ssh_opts],
         "user_droplet_ssh_command" => options[:ssh_command],
-        "user_quiet" => options[:quiet]
+        "user_droplet_ssh_wait"  => options[:wait],
+        "user_quiet" => options[:quiet],
       })
     end
 

lib/tugboat/middleware/base.rb

@@ -39,6 +39,29 @@ def verify_credentials(ocean, say_success=false)
         say "Authentication with DigitalOcean was successful.", :green if say_success
       end
 
+      def wait_for_state(droplet_id, desired_state,ocean)
+        start_time = Time.now
+
+        response = ocean.droplet.show droplet_id
+
+        say ".", nil, false
+
+        if !response.success?
+          say "Failed to get status of Droplet: #{response.message}", :red
+          exit 1
+        end
+
+        while response.droplet.status != desired_state do
+          sleep 2
+          response = ocean.droplet.show droplet_id
+          say ".", nil, false
+        end
+
+        total_time = (Time.now - start_time).to_i
+
+        say "done#{CLEAR} (#{total_time}s)", :green
+      end
+
       # Get all pages of droplets
       def get_droplet_list(ocean)
 

lib/tugboat/middleware/ssh_droplet.rb

@@ -49,6 +49,12 @@ def call(env)
         end
 
         host_string = "#{ssh_user}@#{host_ip}"
+
+        if env['user_droplet_ssh_wait']
+          say "Wait flag given, waiting for droplet to become active"
+          wait_for_state(env["droplet_id"],'active',env['barge'])
+        end
+
         say "Attempting SSH: #{host_string}"
 
         options << host_string

lib/tugboat/middleware/wait_for_state.rb

@@ -6,26 +6,7 @@ def call(env)
 
         say "Waiting for droplet to become #{env["user_droplet_desired_state"]}.", nil, false
 
-        start_time = Time.now
-
-        response = ocean.droplet.show env["droplet_id"]
-
-        say ".", nil, false
-
-        if !response.success?
-          say "Failed to get status of Droplet: #{response.message}", :red
-          exit 1
-        end
-
-        while response.droplet.status != env["user_droplet_desired_state"] do
-          sleep 2
-          response = ocean.droplet.show env["droplet_id"]
-          say ".", nil, false
-        end
-
-        total_time = (Time.now - start_time).to_i
-
-        say "done#{CLEAR} (#{total_time}s)", :green
+        wait_for_state(env["droplet_id"],env["user_droplet_desired_state"],ocean)
 
         @app.call(env)
       end

spec/cli/ssh_cli_spec.rb

@@ -16,6 +16,33 @@
       @cli.ssh("example.com")
     end
 
+    it "wait's until droplet active if -w command is given" do
+      stub_request(:get, "https://api.digitalocean.com/v2/droplets?page=1&per_page=1").
+         with(:headers => {'Accept'=>'*/*', 'Accept-Encoding'=>'gzip;q=1.0,deflate;q=0.6,identity;q=0.3', 'Authorization'=>'Bearer foo', 'Content-Type'=>'application/json', 'User-Agent'=>'Faraday v0.9.2'}).
+         to_return(:status => 200, :body => "", :headers => {})
+
+      stub_request(:get, "https://api.digitalocean.com/v2/droplets/6918990?per_page=200").
+         with(:headers => {'Accept'=>'*/*', 'Accept-Encoding'=>'gzip;q=1.0,deflate;q=0.6,identity;q=0.3', 'Authorization'=>'Bearer foo', 'Content-Type'=>'application/json', 'User-Agent'=>'Faraday v0.9.2'}).
+         to_return(:status => 200, :body => fixture('show_droplet'), :headers => {})
+
+      stub_request(:get, "https://api.digitalocean.com/v2/droplets?page=1&per_page=200").
+        to_return(:headers => {'Content-Type' => 'application/json'}, :status => 200, :body => fixture("show_droplets"))
+      allow(Kernel).to receive(:exec).with('ssh', anything(), anything(),anything(), anything(),anything(), anything(),anything(), anything(),anything(), anything(),anything(), anything(),anything())
+
+      @cli.options = @cli.options.merge(:wait => true)
+
+      @cli.ssh("example.com")
+
+      expect($stdout.string).to eq <<-eos
+Droplet fuzzy name provided. Finding droplet ID...done\e[0m, 6918990 (example.com)
+Executing SSH on Droplet (example.com)...
+Wait flag given, waiting for droplet to become active
+.done\e[0m (0s)
+Attempting SSH: [email protected]
+SShing with options: -o LogLevel=ERROR -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null -o IdentitiesOnly=yes -i #{Dir.home}/.ssh/id_rsa2 -p 33 [email protected]
+      eos
+    end
+
     it "does not allow ssh into a droplet that is inactive" do
       stub_request(:get, "https://api.digitalocean.com/v2/droplets?page=1&per_page=1").
          with(:headers => {'Accept'=>'*/*', 'Accept-Encoding'=>'gzip;q=1.0,deflate;q=0.6,identity;q=0.3', 'Authorization'=>'Bearer foo', 'Content-Type'=>'application/json', 'User-Agent'=>'Faraday v0.9.2'}).