Make the UploadError exception include an array of the hosts that failed

git-svn-id: http://svn.rubyonrails.org/rails/tools/capistrano@6966 5ecf4fe2-1ee6-0310-87b1-e25e094e27de

CHANGELOG

@@ -1,5 +1,7 @@
 *SVN*
 
+* Make the UploadError exception include an array of the hosts that failed [rob@inversepath.com]
+
 * Allow "empty" roles to be declared [Jamis Buck]
 
 * Mercurial SCM module [Tobias Luetke, Matthew Elder]

lib/capistrano/errors.rb

@@ -3,10 +3,13 @@ module Capistrano
 
   class CaptureError < Error; end
   class ConnectionError < Error; end
-  class UploadError < Error; end
   class NoSuchTaskError < Error; end
 
+  class UploadError < Error
+    attr_accessor :hosts
+  end
+
   class CommandError < Error
     attr_accessor :hosts
   end
-end
+end

lib/capistrano/upload.rb

@@ -1,4 +1,5 @@
 require 'net/sftp'
+require 'net/sftp/operations/errors'
 require 'capistrano/errors'
 
 module Capistrano
@@ -62,12 +63,22 @@ module Capistrano
       while running?
         @sftps.each do |sftp|
           next if sftp.channel[:done]
-          sftp.channel.connection.process(true)
+          begin
+            sftp.channel.connection.process(true)
+          rescue Net::SFTP::Operations::StatusException => error
+            logger.important "uploading failed: #{error.description}", sftp.channel[:server] if logger
+            failed!(sftp)
+          end
         end
       end
       logger.trace "upload finished" if logger
 
-      raise UploadError, "upload of #{filename} failed on one or more hosts" if failed > 0
+      if (failed = @sftps.select { |sftp| sftp.channel[:failed] }).any?
+        hosts = failed.map { |sftp| sftp.channel[:server] }
+        error = UploadError.new("upload of #{filename} failed on #{hosts.join(',')}")
+        error.hosts = hosts
+        raise error
+      end
 
       self
     end
@@ -84,7 +95,9 @@ module Capistrano
           sftp = session.sftp
           sftp.connect unless sftp.state == :open
 
+          sftp.channel[:server] = server
           sftp.channel[:done] = false
+          sftp.channel[:failed] = false
           sftp.open(filename, IO::WRONLY | IO::CREAT | IO::TRUNC, options[:mode] || 0660) do |status, handle|
             break unless check_status(sftp, "open #{filename}", server, status)
             
@@ -118,6 +131,7 @@ module Capistrano
       def failed!(sftp)
         completed!(sftp)
         @failed += 1
+        sftp.channel[:failed] = true
       end
 
       def completed!(sftp)

test/upload_test.rb

@@ -16,9 +16,7 @@ class UploadTest < Test::Unit::TestCase
     new_sftp = Proc.new do |state|
       sftp = mock("sftp", :state => state, :open => nil)
       sftp.expects(:connect) unless state == :open
-      chan = mock("channel")
-      chan.expects(:[]=).with(:done, false)
-      sftp.expects(:channel).returns(chan)
+      sftp.stubs(:channel).returns({})
       sftp
     end
 
@@ -34,11 +32,7 @@ class UploadTest < Test::Unit::TestCase
   end
 
   def test_process_when_sftp_open_fails_should_raise_error
-    channel = mock("channel")
-    channel.expects(:[]=).with(:done, false) # initialized to false
-    channel.expects(:[]=).with(:done, true)  # set to true when done
-    sftp = mock("sftp", :state => :open)
-    sftp.expects(:channel).times(2).returns(channel)
+    sftp = mock_sftp
     sftp.expects(:open).with("test.txt", @mode, 0660).yields(mock("status", :code => "bad status", :message => "bad status"), :file_handle)
     session = mock("session", :sftp => sftp, :xserver => server("capistrano"))
     upload = Capistrano::Upload.new([session], "test.txt", :data => "data", :logger => stub_everything)
@@ -48,11 +42,7 @@ class UploadTest < Test::Unit::TestCase
   end
 
   def test_process_when_sftp_write_fails_should_raise_error
-    channel = mock("channel")
-    channel.expects(:[]=).with(:done, false)
-    channel.expects(:[]=).with(:done, true)
-    sftp = mock("sftp", :state => :open)
-    sftp.expects(:channel).times(2).returns(channel)
+    sftp = mock_sftp
     sftp.expects(:open).with("test.txt", @mode, 0660).yields(mock("status1", :code => Net::SFTP::Session::FX_OK), :file_handle)
     sftp.expects(:write).with(:file_handle, "data").yields(mock("status2", :code => "bad status", :message => "bad status"))
     session = mock("session", :sftp => sftp, :xserver => server("capistrano"))
@@ -63,11 +53,7 @@ class UploadTest < Test::Unit::TestCase
   end
 
   def test_process_when_sftp_succeeds_should_raise_nothing
-    channel = mock("channel")
-    channel.expects(:[]=).with(:done, false)
-    channel.expects(:[]=).with(:done, true)
-    sftp = mock("sftp", :state => :open)
-    sftp.expects(:channel).times(2).returns(channel)
+    sftp = mock_sftp
     sftp.expects(:open).with("test.txt", @mode, 0660).yields(mock("status1", :code => Net::SFTP::Session::FX_OK), :file_handle)
     sftp.expects(:write).with(:file_handle, "data").yields(mock("status2", :code => Net::SFTP::Session::FX_OK))
     sftp.expects(:close_handle).with(:file_handle).yields
@@ -81,12 +67,10 @@ class UploadTest < Test::Unit::TestCase
   def test_process_should_loop_while_running
     con = mock("connection")
     con.expects(:process).with(true).times(10)
-    channel = mock("channel")
-    channel.expects(:[]=).with(:done, false)
-    channel.expects(:[]).with(:done).returns(false).times(10)
+    channel = {}
     channel.expects(:connection).returns(con).times(10)
     sftp = mock("sftp", :state => :open, :open => nil)
-    sftp.expects(:channel => channel).times(21)
+    sftp.stubs(:channel).returns(channel)
     session = mock("session", :sftp => sftp, :xserver => server("capistrano"))
     upload = Capistrano::Upload.new([session], "test.txt", :data => "data")
     upload.expects(:running?).times(11).returns(*([true]*10 + [false]))
@@ -95,23 +79,37 @@ class UploadTest < Test::Unit::TestCase
 
   def test_process_should_loop_but_not_process_done_channels
     new_sftp = Proc.new do |done|
+      channel = {}
+      channel[:needs_done] = done
+
       if !done
         con = mock("connection")
         con.expects(:process).with(true).times(10)
+        channel.expects(:connection).returns(con).times(10)
       end
-      channel = mock("channel")
-      channel.expects(:[]=).with(:done, false)
-      channel.expects(:[]).with(:done).returns(done).times(10)
-      channel.expects(:connection).returns(con).times(10) if !done
+
       sftp = mock("sftp", :state => :open, :open => nil)
-      sftp.expects(:channel => channel).times(done ? 11 : 21)
+      sftp.stubs(:channel).returns(channel)
       sftp
     end
 
-    sessions = [mock("session", :sftp => new_sftp[true], :xserver => server("capistrano")),
-                mock("session", :sftp => new_sftp[false], :xserver => server("cap2"))]
+    sessions = [stub("session", :sftp => new_sftp[true], :xserver => server("capistrano")),
+                stub("session", :sftp => new_sftp[false], :xserver => server("cap2"))]
     upload = Capistrano::Upload.new(sessions, "test.txt", :data => "data")
+
+    # make sure the sftp channels we wanted to be done, start as done
+    # (Upload.new marks each channel as not-done, so we have to do it here)
+    sessions.each { |s| s.sftp.channel[:done] = true if s.sftp.channel[:needs_done] }
     upload.expects(:running?).times(11).returns(*([true]*10 + [false]))
     upload.process!
   end
+
+  private
+
+    def mock_sftp
+      sftp = mock("sftp", :state => :open)
+      sftp.stubs(:channel).returns(Hash.new)
+      yield sftp if block_given?
+      sftp
+    end
 end