allow session#use to declare multiple servers in a single call

2 files changed, 28 insertions, 18 deletions

diff --git a/lib/net/ssh/multi/session.rb b/lib/net/ssh/multi/session.rb
index c7462d6..0f9df0b 100644
--- a/lib/net/ssh/multi/session.rb
+++ b/lib/net/ssh/multi/session.rb
@@ -181,16 +181,28 @@ module Net; module SSH; module Multi
     #   session.use 'host'
     #   session.use 'user@host2', :via => nil
     #   session.use 'host3', :user => "user3", :via => Net::SSH::Gateway.new('gateway.host', 'user')
-    def use(host, options={})
-      server = Server.new(self, host, {:via => default_gateway}.merge(options))
-      exists = servers.index(server)
-      if exists
-        server = servers[exists]
-      else
-        servers << server
-        group [] => server
+    #
+    # If only a single host is given, the new server instance is returned. You
+    # can give multiple hosts at a time, though, in which case an array of
+    # server instances will be returned.
+    #
+    #   server1, server2 = session.use "host1", "host2"
+    def use(*hosts)
+      options = hosts.last.is_a?(Hash) ? hosts.pop : {}
+
+      results = hosts.map do |host|
+        server = Server.new(self, host, {:via => default_gateway}.merge(options))
+        exists = servers.index(server)
+        if exists
+          server = servers[exists]
+        else
+          servers << server
+        end
+        server
       end
-      server
+
+      group [] => results
+      results.length > 1 ? results : results.first
     end
 
     # Returns the set of servers that match the given criteria. It can be used
diff --git a/test/session_test.rb b/test/session_test.rb
index 6446ace..5a948d8 100644
--- a/test/session_test.rb
+++ b/test/session_test.rb
@@ -67,8 +67,7 @@ class SessionTest < Test::Unit::TestCase
   end
 
   def test_use_with_duplicate_server_will_not_add_server_twice
-    s1 = @session.use('host')
-    s2 = @session.use('host')
+    s1, s2 = @session.use('host', 'host')
     assert_equal 1, @session.servers.length
     assert_equal s1.object_id, s2.object_id
   end
@@ -96,8 +95,7 @@ class SessionTest < Test::Unit::TestCase
   end
 
   def test_on_should_return_subsession_containing_only_the_given_servers
-    s1 = @session.use('h1')
-    s2 = @session.use('h2')
+    s1, s2 = @session.use('h1', 'h2')
     subsession = @session.on(s1, s2)
     assert_equal [s1, s2], subsession.servers
   end
@@ -113,18 +111,18 @@ class SessionTest < Test::Unit::TestCase
   end
 
   def test_servers_for_should_return_all_servers_if_no_arguments
-    srv1, srv2, srv3 = @session.use('h1'), @session.use('h2'), @session.use('h3')
+    srv1, srv2, srv3 = @session.use('h1', 'h2', 'h3')
     assert_equal [srv1, srv2, srv3], @session.servers_for.sort
   end
 
   def test_servers_for_should_return_servers_only_for_given_group
-    srv1, srv2, srv3 = @session.use('h1'), @session.use('h2'), @session.use('h3')
+    srv1, srv2, srv3 = @session.use('h1', 'h2', 'h3')
     @session.group :app => [srv1, srv2], :db => [srv3]
     assert_equal [srv1, srv2], @session.servers_for(:app).sort
   end
 
   def test_servers_for_should_not_return_duplicate_servers
-    srv1, srv2, srv3 = @session.use('h1'), @session.use('h2'), @session.use('h3')
+    srv1, srv2, srv3 = @session.use('h1', 'h2', 'h3')
     @session.group :app => [srv1, srv2], :db => [srv2, srv3]
     assert_equal [srv1, srv2, srv3], @session.servers_for(:app, :db).sort
   end
@@ -136,7 +134,7 @@ class SessionTest < Test::Unit::TestCase
   end
 
   def test_close_should_close_server_sessions
-    srv1, srv2 = @session.use('h1'), @session.use('h2')
+    srv1, srv2 = @session.use('h1', 'h2')
     srv1.expects(:close_channels)
     srv2.expects(:close_channels)
     srv1.expects(:close)
@@ -189,7 +187,7 @@ class SessionTest < Test::Unit::TestCase
 
   def test_process_should_call_select_on_combined_readers_and_writers_from_all_servers
     @session.expects(:postprocess).with([:b, :c], [:a, :c])
-    srv1, srv2, srv3 = @session.use('h1'), @session.use('h2'), @session.use('h3')
+    srv1, srv2, srv3 = @session.use('h1', 'h2', 'h3')
     srv1.expects(:readers).returns([:a])
     srv1.expects(:writers).returns([:a])
     srv2.expects(:readers).returns([])