test/configuration/connections_test.rb in capistrano-2.1.0 vs test/configuration/connections_test.rb in capistrano-2.2.0

@@ -215,11 +215,10 @@
   def test_execute_on_servers_should_not_try_to_connect_to_hosts_with_connection_errors_with_on_errors_continue
     list = [server("cap1"), server("cap2")]
     @config.current_task = mock_task(:on_error => :continue)
     @config.expects(:find_servers_for_task).with(@config.current_task, {}).returns(list)
     Capistrano::SSH.expects(:connect).times(2).raises(Exception).then.returns(:success)
-    @config.expects(:failed!).with(server("cap1"))
     @config.execute_on_servers do |servers|
       assert_equal %w(cap2), servers.map { |s| s.host }
     end
     @config.expects(:find_servers_for_task).with(@config.current_task, {}).returns(list)
     @config.execute_on_servers do |servers|
@@ -258,20 +257,56 @@
     @config.expects(:find_servers_for_task).with(@config.current_task, {}).returns([cap1, cap2])
     @config.execute_on_servers do |servers|
       assert_equal %w(cap2), servers.map { |s| s.host }
     end
   end
-
+  
   def test_connect_should_establish_connections_to_all_servers_in_scope
     assert @config.sessions.empty?
     @config.current_task = mock_task
     @config.expects(:find_servers_for_task).with(@config.current_task, {}).returns([server("cap1"), server("cap2"), server("cap3")])
     Capistrano::SSH.expects(:connect).times(3).returns(:success)
     @config.connect!
     assert_equal %w(cap1 cap2 cap3), @config.sessions.keys.sort.map { |s| s.host }
   end
-
+  
+  def test_execute_on_servers_should_only_run_on_tasks_max_hosts_hosts_at_once
+    cap1 = server("cap1")
+    cap2 = server("cap2")
+    connection1 = mock()
+    connection2 = mock()
+    connection1.expects(:close)
+    connection2.expects(:close)
+    @config.current_task = mock_task(:max_hosts => 1)
+    @config.expects(:find_servers_for_task).with(@config.current_task, {}).returns([cap1, cap2])
+    Capistrano::SSH.expects(:connect).times(2).returns(connection1).then.returns(connection2)
+    block_called = 0
+    @config.execute_on_servers do |servers|
+      block_called += 1
+      assert_equal 1, servers.size
+    end
+    assert_equal 2, block_called
+  end
+  
+  def test_execute_on_servers_should_only_run_on_max_hosts_hosts_at_once
+    cap1 = server("cap1")
+    cap2 = server("cap2")
+    connection1 = mock()
+    connection2 = mock()
+    connection1.expects(:close)
+    connection2.expects(:close)
+    @config.current_task = mock_task(:max_hosts => 1)
+    @config.expects(:find_servers_for_task).with(@config.current_task, {}).returns([cap1, cap2])
+    Capistrano::SSH.expects(:connect).times(2).returns(connection1).then.returns(connection2)
+    block_called = 0
+    @config.execute_on_servers do |servers|
+      block_called += 1
+      assert_equal 1, servers.size
+    end
+    assert_equal 2, block_called
+  end
+  
   def test_connect_should_honor_once_option
     assert @config.sessions.empty?
     @config.current_task = mock_task
     @config.expects(:find_servers_for_task).with(@config.current_task, :once => true).returns([server("cap1"), server("cap2"), server("cap3")])
     Capistrano::SSH.expects(:connect).returns(:success)
@@ -281,8 +316,13 @@
 
   private
 
     def mock_task(options={})
       continue_on_error = options[:on_error] == :continue
-      stub("task", :fully_qualified_name => "name", :options => options, :continue_on_error? => continue_on_error)
+      stub("task",
+        :fully_qualified_name => "name",
+        :options => options,
+        :continue_on_error? => continue_on_error,
+        :max_hosts => options[:max_hosts]
+      )
     end
 end
\ No newline at end of file