make with() and on() yield a new subsession object that encapsulates the set of matching servers, and add a new servers_for method.

7 files changed, 430 insertions, 380 deletions

diff --git a/README.rdoc b/README.rdoc
index 0818e47..3a90d57 100644
--- a/README.rdoc
+++ b/README.rdoc
@@ -33,11 +33,11 @@ In a nutshell:
       session.use 'app2', 'user'
     end
 
-    # execute commands
+    # execute commands on all servers
     session.exec "uptime"
 
     # execute commands on a subset of servers
-    session.with(:app) { session.exec "hostname" }
+    session.with(:app).exec "hostname"
 
     # run the aggregated event loop
     session.loop
diff --git a/lib/net/ssh/multi.rb b/lib/net/ssh/multi.rb
index ae24f99..8f393df 100644
--- a/lib/net/ssh/multi.rb
+++ b/lib/net/ssh/multi.rb
@@ -28,11 +28,11 @@ module Net; module SSH
   #       session.use 'app2', 'user'
   #     end
   # 
-  #     # execute commands
+  #     # execute commands on all servers
   #     session.exec "uptime"
   # 
   #     # execute commands on a subset of servers
-  #     session.with(:app) { session.exec "hostname" }
+  #     session.with(:app).exec "hostname"
   # 
   #     # run the aggregated event loop
   #     session.loop
diff --git a/lib/net/ssh/multi/session.rb b/lib/net/ssh/multi/session.rb
index 99819e1..561f9d1 100644
--- a/lib/net/ssh/multi/session.rb
+++ b/lib/net/ssh/multi/session.rb
@@ -3,6 +3,8 @@ require 'net/ssh/gateway'
 require 'net/ssh/multi/server'
 require 'net/ssh/multi/channel'
 require 'net/ssh/multi/pending_connection'
+require 'net/ssh/multi/session_actions'
+require 'net/ssh/multi/subsession'
 
 module Net; module SSH; module Multi
   # Represents a collection of connections to various servers. It provides an
@@ -27,11 +29,11 @@ module Net; module SSH; module Multi
   #       session.use 'app2', 'user'
   #     end
   # 
-  #     # execute commands
+  #     # execute commands on all servers
   #     session.exec "uptime"
   # 
   #     # execute commands on a subset of servers
-  #     session.with(:app) { session.exec "hostname" }
+  #     session.with(:app).exec "hostname"
   # 
   #     # run the aggregated event loop
   #     session.loop
@@ -41,6 +43,8 @@ module Net; module SSH; module Multi
   # You can force the connections to be opened immediately, though, using the
   # #connect! method.
   class Session
+    include SessionActions
+
     # The list of Net::SSH::Multi::Server definitions managed by this session.
     attr_reader :servers
 
@@ -63,11 +67,6 @@ module Net; module SSH; module Multi
     # See #use and #group.
     attr_reader :open_groups #:nodoc:
 
-    # The list of "active" groups, which will be used to restrict subsequent
-    # commands. This is actually a Hash, mapping group names to their corresponding
-    # constraints (see #with).
-    attr_reader :active_groups #:nodoc:
-
     # Creates a new Net::SSH::Multi::Session instance. Initially, it contains
     # no server definitions, no group definitions, and no default gateway.
     #
@@ -85,7 +84,6 @@ module Net; module SSH; module Multi
       @servers = []
       @groups = {}
       @gateway = nil
-      @active_groups = {}
       @open_groups = []
       @connect_threads = []
 
@@ -182,24 +180,22 @@ module Net; module SSH; module Multi
       server
     end
 
-    # Restricts the set of servers that will be targeted by commands within
-    # the associated block. It can be used in either of two ways (or both ways
-    # used together).
+    # Returns the set of servers that match the given criteria. It can be used
+    # in any (or all) of three ways.
     #
-    # First, you can simply specify a list of group names. All servers in all
-    # named groups will be the target of the commands. (Nested calls to #with
-    # are cumulative.)
+    # First, you can omit any arguments. In this case, the full list of servers
+    # will be returned.
     #
-    #   # execute 'hostname' on all servers in the :app group, and 'uptime'
-    #   # on all servers in either :app or :db.
-    #   session.with(:app) do
-    #     session.exec('hostname')
-    #     session.with(:db) do
-    #       session.exec('uptime')
-    #     end
-    #   end
+    #   all = session.servers_for
+    #
+    # Second, you can simply specify a list of group names. All servers in all
+    # named groups will be returned. If a server belongs to multiple matching
+    # groups, then it will appear only once in the list (the resulting list
+    # will contain only unique servers).
     #
-    # Secondly, you can specify a hash with group names as keys, and property
+    #   servers = session.servers_for(:app, :db)
+    #
+    # Last, you can specify a hash with group names as keys, and property
     # constraints as the values. These property constraints are either "only"
     # constraints (which restrict the set of servers to "only" those that match
     # the given properties) or "except" constraints (which restrict the set of
@@ -212,95 +208,76 @@ module Net; module SSH; module Multi
     #     session.use 'dbslve2', 'user2'
     #   end
     #
-    #   # execute the given rake task ONLY on the servers in the :db group
-    #   # which have the :primary property set to true.
-    #   session.with :db => { :only => { :primary => true } } do
-    #     session.exec "rake db:migrate"
-    #   end
+    #   # return ONLY on the servers in the :db group which have the :primary
+    #   # property set to true.
+    #   primary = session.servers_for(:db => { :only => { :primary => true } })
     #
     # You can, naturally, combine these methods:
     #
     #   # all servers in :app and :web, and all servers in :db with the
     #   # :primary property set to true
-    #   session.with :app, :web, :db => { :only => { :primary => true } } do
-    #     # ...
-    #   end
-    def with(*groups)
-      saved_groups = active_groups.dup
-
-      new_map = groups.inject({}) do |map, group|
-        if group.is_a?(Hash)
-          group.each do |gr, value|
-            raise ArgumentError, "the value for any group must be a Hash" unless value.is_a?(Hash)
-            bad_keys = value.keys - [:only, :except]
-            raise ArgumentError, "unknown constraint(s): #{bad_keys.inspect}" unless bad_keys.empty?
-            map[gr] = (active_groups[gr] || {}).merge(value)
+    #   servers = session.servers_for(:app, :web, :db => { :only => { :primary => true } })
+    def servers_for(*criteria)
+      if criteria.empty?
+        servers
+      else
+        # normalize the criteria list, so that every entry is a key to a
+        # criteria hash (possibly empty).
+        criteria = criteria.inject({}) do |hash, entry|
+          case entry
+          when Hash then hash.merge(entry)
+          else hash.merge(entry => {})
           end
-        else
-          map[group] = active_groups[group] || {}
         end
-        map
-      end
-
-      active_groups.update(new_map)
-      yield self
-    ensure
-      active_groups.replace(saved_groups)
-    end
 
-    # Works as #with, but for specific servers rather than groups. In other
-    # words, you can use this to restrict actions within the block to only
-    # a specific list of servers. It works by creating an ad-hoc group, adding
-    # the servers to that group, and then making that group the only active
-    # group. (Note that because of this, you cannot nest #on within #with,
-    # though you could nest #with inside of #on.)
-    #
-    #   srv = session.use('host', 'user')
-    #   # ...
-    #   session.on(srv) do
-    #     session.exec('hostname')
-    #   end
-    def on(*servers)
-      adhoc_group = "adhoc_group_#{servers.hash}_#{rand(0xffffffff)}".to_sym
-      group(adhoc_group => servers)
-      saved_groups = active_groups.dup
-      active_groups.replace(adhoc_group => {})
-      yield self
-    ensure
-      active_groups.replace(saved_groups) if saved_groups
-      groups.delete(adhoc_group)
-    end
+        list = criteria.inject([]) do |server_list, (group, properties)|
+          raise ArgumentError, "the value for any group must be a Hash, but got a #{properties.class} for #{group.inspect}" unless properties.is_a?(Hash)
+          bad_keys = properties.keys - [:only, :except]
+          raise ArgumentError, "unknown constraint(s) #{bad_keys.inspect} for #{group.inspect}" unless bad_keys.empty?
 
-    # Returns the list of Net::SSH sessions for all servers that match the
-    # current scope (e.g., the groups or servers named in the outer #with or
-    # #on calls). If any servers have not yet been connected to, this will
-    # block until the connections have been made.
-    def active_sessions
-      list = if active_groups.empty?
-        servers
-      else
-        active_groups.inject([]) do |list, (group, properties)|
-          servers = groups[group].select do |server|
+          servers = (groups[group] || []).select do |server|
             (properties[:only] || {}).all? { |prop, value| server[prop] == value } &&
             !(properties[:except] || {}).any? { |prop, value| server[prop] == value }
           end
-          list.concat(servers)
+          server_list.concat(servers)
         end
-      end
 
-      list.uniq!
-      threads = list.map { |server| Thread.new { server.session(true) } if server.session.nil? }
-      threads.each { |thread| thread.join if thread }
+        list.uniq
+      end
+    end
 
-      list.map { |server| server.session }
+    # Returns a new Net::SSH::Multi::Subsession instance consisting of the
+    # servers that meet the given criteria. If a block is given, the
+    # subsession will be yielded to it. See #servers_for for a discussion of
+    # how these criteria are interpreted.
+    #
+    #   session.with(:app).exec('hostname')
+    #
+    #   session.with(:app, :db => { :primary => true }) do |s|
+    #     s.exec 'date'
+    #     s.exec 'uptime'
+    #   end
+    def with(*groups)
+      subsession = Subsession.new(self, servers_for(*groups))
+      yield subsession if block_given?
+      subsession
     end
 
-    # Connections are normally established lazily, as soon as they are needed.
-    # This method forces all servers selected by the current scope to connect,
-    # if they have not yet been connected.
-    def connect!
-      active_sessions
-      self
+    # Works as #with, but for specific servers rather than groups. It will
+    # return a new subsession (Net::SSH::Multi::Subsession) consisting of
+    # the given servers. (Note that it requires that the servers in question
+    # have been created via calls to #use on this session object, or things
+    # will not work quite right.) If a block is given, the new subsession
+    # will also be yielded to the block.
+    #
+    #   srv1 = session.use('host1', 'user')
+    #   srv2 = session.use('host2', 'user')
+    #   # ...
+    #   session.on(srv1, srv2).exec('hostname')
+    def on(*servers)
+      subsession = Subsession.new(self, servers)
+      yield subsession if block_given?
+      subsession
     end
 
     # Closes the multi-session by shutting down all open server sessions, and
@@ -314,14 +291,6 @@ module Net; module SSH; module Multi
       default_gateway.shutdown! if default_gateway
     end
 
-    # Returns +true+ if any server has an open SSH session that is currently
-    # processing any channels. If +include_invisible+ is +false+ (the default)
-    # then invisible channels (such as those created by port forwarding) will
-    # not be counted; otherwise, they will be.
-    def busy?(include_invisible=false)
-      servers.any? { |server| server.busy?(include_invisible) }
-    end
-
     alias :loop_forever :loop
 
     # Run the aggregated event loop for all open server sessions, until the given
@@ -357,116 +326,6 @@ module Net; module SSH; module Multi
       end
     end
 
-    # Sends a global request to all active sessions (see #active_sessions).
-    # This can be used to (e.g.) ping the remote servers to prevent them from
-    # timing out.
-    #
-    #   session.send_global_request("keep-alive@openssh.com")
-    #
-    # If a block is given, it will be invoked when the server responds, with
-    # two arguments: the Net::SSH connection that is responding, and a boolean
-    # indicating whether the request succeeded or not.
-    def send_global_request(type, *extra, &callback)
-      active_sessions.each { |ssh| ssh.send_global_request(type, *extra, &callback) }
-      self
-    end
-
-    # Asks all active sessions (see #active_sessions) to open a new channel.
-    # When each server responds, the +on_confirm+ block will be invoked with
-    # a single argument, the channel object for that server. This means that
-    # the block will be invoked one time for each active session.
-    #
-    # All new channels will be collected and returned, aggregated into a new
-    # Net::SSH::Multi::Channel instance.
-    #
-    # Note that the channels are "enhanced" slightly--they have two properties
-    # set on them automatically, to make dealing with them in a multi-session
-    # environment slightly easier:
-    #
-    # * :server => the Net::SSH::Multi::Server instance that spawned the channel
-    # * :host => the host name of the server
-    #
-    # Having access to these things lets you more easily report which host
-    # (e.g.) data was received from:
-    #
-    #   session.open_channel do |channel|
-    #     channel.exec "command" do |ch, success|
-    #       ch.on_data do |ch, data|
-    #         puts "got data #{data} from #{ch[:host]}"
-    #       end
-    #     end
-    #   end
-    def open_channel(type="session", *extra, &on_confirm)
-      channels = active_sessions.map do |ssh|
-        ssh.open_channel(type, *extra) do |c|
-          c[:server] = c.connection[:server]
-          c[:host] = c.connection[:server].host
-          on_confirm[c] if on_confirm
-        end
-      end
-      Multi::Channel.new(self, channels)
-    end
-
-    # A convenience method for executing a command on multiple hosts and
-    # either displaying or capturing the output. It opens a channel on all
-    # active sessions (see #open_channel and #active_sessions), and then
-    # executes a command on each channel (Net::SSH::Connection::Channel#exec).
-    #
-    # If a block is given, it will be invoked whenever data is received across
-    # the channel, with three arguments: the channel object, a symbol identifying
-    # which output stream the data was received on (+:stdout+ or +:stderr+)
-    # and a string containing the data that was received:
-    #
-    #   session.exec("command") do |ch, stream, data|
-    #     puts "[#{ch[:host]} : #{stream}] #{data}"
-    #   end
-    #
-    # If no block is given, all output will be written to +$stdout+ or
-    # +$stderr+, as appropriate.
-    #
-    # Note that #exec will also capture the exit status of the process in the
-    # +:exit_status+ property of each channel. Since #exec returns all of the
-    # channels in a Net::SSH::Multi::Channel object, you can check for the
-    # exit status like this:
-    #
-    #   channel = session.exec("command") { ... }
-    #   channel.wait
-    #
-    #   if channel.any? { |c| c[:exit_status] != 0 }
-    #     puts "executing failed on at least one host!"
-    #   end
-    def exec(command, &block)
-      open_channel do |channel|
-        channel.exec(command) do |ch, success|
-          raise "could not execute command: #{command.inspect} (#{ch[:host]})" unless success
-
-          channel.on_data do |ch, data|
-            if block
-              block.call(ch, :stdout, data)
-            else
-              data.chomp.each_line do |line|
-                $stdout.puts("[#{ch[:host]}] #{line}")
-              end
-            end
-          end
-
-          channel.on_extended_data do |ch, type, data|
-            if block
-              block.call(ch, :stderr, data)
-            else
-              data.chomp.each_line do |line|
-                $stderr.puts("[#{ch[:host]}] #{line}")
-              end
-            end
-          end
-
-          channel.on_request("exit-status") do |ch, data|
-            ch[:exit_status] = data.read_long
-          end
-        end
-      end
-    end
-
     # Runs the preprocess stage on all servers. Returns false if the block
     # returns false, and true if there either is no block, or it returns true.
     # This is called as part of the #process method.
diff --git a/lib/net/ssh/multi/session_actions.rb b/lib/net/ssh/multi/session_actions.rb
new file mode 100644
index 0000000..35c69ed
--- /dev/null
+++ b/lib/net/ssh/multi/session_actions.rb
@@ -0,0 +1,153 @@
+module Net; module SSH; module Multi
+
+  # This module represents the actions that are available on session
+  # collections. Any class that includes this module needs only provide a
+  # +servers+ method that returns a list of Net::SSH::Multi::Server
+  # instances, and the rest just works. See Net::SSH::Multi::Session and
+  # Net::SSH::Multi::Subsession for consumers of this module.
+  module SessionActions
+    # Returns the session that is the "master". This defaults to +self+, but
+    # classes that include this module may wish to change this if they are
+    # subsessions that depend on a master session.
+    def master
+      self
+    end
+
+    # Connections are normally established lazily, as soon as they are needed.
+    # This method forces all servers in the current container to have their
+    # connections established immediately, blocking until the connections have
+    # been made.
+    def connect!
+      sessions
+      self
+    end
+
+    # Returns +true+ if any server in the current container has an open SSH
+    # session that is currently processing any channels. If +include_invisible+
+    # is +false+ (the default) then invisible channels (such as those created
+    # by port forwarding) will not be counted; otherwise, they will be.
+    def busy?(include_invisible=false)
+      servers.any? { |server| server.busy?(include_invisible) }
+    end
+
+    # Returns an array of all SSH sessions, blocking until all sessions have
+    # connected.
+    def sessions
+      threads = servers.map { |server| Thread.new { server.session(true) } if server.session.nil? }
+      threads.each { |thread| thread.join if thread }
+      servers.map { |server| server.session }
+    end
+
+    # Sends a global request to the sessions for all contained servers
+    # (see #sessions). This can be used to (e.g.) ping the remote servers to
+    # prevent them from timing out.
+    #
+    #   session.send_global_request("keep-alive@openssh.com")
+    #
+    # If a block is given, it will be invoked when the server responds, with
+    # two arguments: the Net::SSH connection that is responding, and a boolean
+    # indicating whether the request succeeded or not.
+    def send_global_request(type, *extra, &callback)
+      sessions.each { |ssh| ssh.send_global_request(type, *extra, &callback) }
+      self
+    end
+
+    # Asks all sessions for all contained servers (see #sessions) to open a
+    # new channel. When each server responds, the +on_confirm+ block will be
+    # invoked with a single argument, the channel object for that server. This
+    # means that the block will be invoked one time for each session.
+    #
+    # All new channels will be collected and returned, aggregated into a new
+    # Net::SSH::Multi::Channel instance.
+    #
+    # Note that the channels are "enhanced" slightly--they have two properties
+    # set on them automatically, to make dealing with them in a multi-session
+    # environment slightly easier:
+    #
+    # * :server => the Net::SSH::Multi::Server instance that spawned the channel
+    # * :host => the host name of the server
+    #
+    # Having access to these things lets you more easily report which host
+    # (e.g.) data was received from:
+    #
+    #   session.open_channel do |channel|
+    #     channel.exec "command" do |ch, success|
+    #       ch.on_data do |ch, data|
+    #         puts "got data #{data} from #{ch[:host]}"
+    #       end
+    #     end
+    #   end
+    def open_channel(type="session", *extra, &on_confirm)
+      channels = sessions.map do |ssh|
+        ssh.open_channel(type, *extra) do |c|
+          c[:server] = c.connection[:server]
+          c[:host] = c.connection[:server].host
+          on_confirm[c] if on_confirm
+        end
+      end
+      Multi::Channel.new(master, channels)
+    end
+
+    # A convenience method for executing a command on multiple hosts and
+    # either displaying or capturing the output. It opens a channel on all
+    # active sessions (see #open_channel and #active_sessions), and then
+    # executes a command on each channel (Net::SSH::Connection::Channel#exec).
+    #
+    # If a block is given, it will be invoked whenever data is received across
+    # the channel, with three arguments: the channel object, a symbol identifying
+    # which output stream the data was received on (+:stdout+ or +:stderr+)
+    # and a string containing the data that was received:
+    #
+    #   session.exec("command") do |ch, stream, data|
+    #     puts "[#{ch[:host]} : #{stream}] #{data}"
+    #   end
+    #
+    # If no block is given, all output will be written to +$stdout+ or
+    # +$stderr+, as appropriate.
+    #
+    # Note that #exec will also capture the exit status of the process in the
+    # +:exit_status+ property of each channel. Since #exec returns all of the
+    # channels in a Net::SSH::Multi::Channel object, you can check for the
+    # exit status like this:
+    #
+    #   channel = session.exec("command") { ... }
+    #   channel.wait
+    #
+    #   if channel.any? { |c| c[:exit_status] != 0 }
+    #     puts "executing failed on at least one host!"
+    #   end
+    def exec(command, &block)
+      open_channel do |channel|
+        channel.exec(command) do |ch, success|
+          raise "could not execute command: #{command.inspect} (#{ch[:host]})" unless success
+
+          channel.on_data do |ch, data|
+            if block
+              block.call(ch, :stdout, data)
+            else
+              data.chomp.each_line do |line|
+                $stdout.puts("[#{ch[:host]}] #{line}")
+              end
+            end
+          end
+
+          channel.on_extended_data do |ch, type, data|
+            if block
+              block.call(ch, :stderr, data)
+            else
+              data.chomp.each_line do |line|
+                $stderr.puts("[#{ch[:host]}] #{line}")
+              end
+            end
+          end
+
+          channel.on_request("exit-status") do |ch, data|
+            ch[:exit_status] = data.read_long
+          end
+        end
+      end
+    end
+
+  end
+
+end; end; end
\ No newline at end of file
diff --git a/lib/net/ssh/multi/subsession.rb b/lib/net/ssh/multi/subsession.rb
new file mode 100644
index 0000000..29e7866
--- /dev/null
+++ b/lib/net/ssh/multi/subsession.rb
@@ -0,0 +1,48 @@
+require 'net/ssh/multi/session_actions'
+
+module Net; module SSH; module Multi
+
+  # A trivial class for representing a subset of servers. It is used
+  # internally for restricting operations to a subset of all defined
+  # servers.
+  #
+  #   subsession = session.with(:app)
+  #   subsession.exec("hostname")
+  class Subsession
+    include SessionActions
+
+    # The master session that spawned this subsession.
+    attr_reader :master
+
+    # The list of servers that this subsession can operate on.
+    attr_reader :servers
+
+    # Create a new subsession of the given +master+ session, that operates
+    # on the given +server_list+.
+    def initialize(master, server_list)
+      @master  = master
+      @servers = server_list.uniq
+    end
+
+    # Works as Array#slice, but returns a new subsession consisting of the
+    # given slice of servers in this subsession. The new subsession will have
+    # the same #master session as this subsession does.
+    #
+    #   s1 = subsession.slice(0)
+    #   s2 = subsession.slice(3, -1)
+    #   s3 = subsession.slice(1..4)
+    def slice(*args)
+      Subsession.new(master, Array(servers.slice(*args)))
+    end
+
+    # Returns a new subsession that consists of only the first server in the
+    # server list of the current subsession. This is just convenience for
+    # #slice(0):
+    #
+    #   s1 = subsession.first
+    def first
+      slice(0)
+    end
+  end
+
+end; end; end
\ No newline at end of file
diff --git a/test/session_actions_test.rb b/test/session_actions_test.rb
new file mode 100644
index 0000000..c4f3e58
--- /dev/null
+++ b/test/session_actions_test.rb
@@ -0,0 +1,124 @@
+require 'common'
+require 'net/ssh/multi/server'
+require 'net/ssh/multi/session_actions'
+
+class SessionActionsTest < Test::Unit::TestCase
+  class SessionActionsContainer
+    include Net::SSH::Multi::SessionActions
+
+    attr_reader :servers
+
+    def initialize
+      @servers = []
+    end
+
+    def use(h, u, o={})
+      server = Net::SSH::Multi::Server.new(self, h, u, o)
+      servers << server
+      server
+    end
+  end
+
+  def setup
+    @session = SessionActionsContainer.new
+  end
+
+  def test_busy_should_be_true_if_any_server_is_busy
+    srv1, srv2, srv3 = @session.use('h1', 'u1'), @session.use('h2', 'u2'), @session.use('h3', 'u3')
+    srv1.stubs(:busy?).returns(false)
+    srv2.stubs(:busy?).returns(false)
+    srv3.stubs(:busy?).returns(true)
+    assert @session.busy?
+  end
+
+  def test_busy_should_be_false_if_all_servers_are_not_busy
+    srv1, srv2, srv3 = @session.use('h1', 'u1', :properties => {:a => 1}), @session.use('h2', 'u2', :properties => {:a => 1, :b => 2}), @session.use('h3', 'u3')
+    srv1.stubs(:busy?).returns(false)
+    srv2.stubs(:busy?).returns(false)
+    srv3.stubs(:busy?).returns(false)
+    assert !@session.busy?
+  end
+
+  def test_send_global_request_should_delegate_to_sessions
+    s1 = mock('ssh')
+    s2 = mock('ssh')
+    s1.expects(:send_global_request).with("a", "b", "c").yields
+    s2.expects(:send_global_request).with("a", "b", "c").yields
+    @session.expects(:sessions).returns([s1, s2])
+    calls = 0
+    @session.send_global_request("a", "b", "c") { calls += 1 }
+    assert_equal 2, calls
+  end
+
+  def test_open_channel_should_delegate_to_sessions_and_set_accessors_on_each_channel_and_return_multi_channel
+    srv1 = @session.use('h1', 'u1')
+    srv2 = @session.use('h2', 'u2')
+    s1 = { :server => srv1 }
+    s2 = { :server => srv2 }
+    c1 = { :stub => :value }
+    c2 = {}
+    c1.stubs(:connection).returns(s1)
+    c2.stubs(:connection).returns(s2)
+    @session.expects(:sessions).returns([s1, s2])
+    s1.expects(:open_channel).with("session").yields(c1).returns(c1)
+    s2.expects(:open_channel).with("session").yields(c2).returns(c2)
+    results = []
+    channel = @session.open_channel do |c|
+      results << c
+    end
+    assert_equal [c1, c2], results
+    assert_equal "h1", c1[:host]
+    assert_equal "h2", c2[:host]
+    assert_equal srv1, c1[:server]
+    assert_equal srv2, c2[:server]
+    assert_instance_of Net::SSH::Multi::Channel, channel
+    assert_equal [c1, c2], channel.channels
+  end
+
+  def test_exec_should_raise_exception_if_channel_cannot_exec_command
+    c = { :host => "host" }
+    @session.expects(:open_channel).yields(c).returns(c)
+    c.expects(:exec).with('something').yields(c, false)
+    assert_raises(RuntimeError) { @session.exec("something") }
+  end
+
+  def test_exec_with_block_should_pass_data_and_extended_data_to_block
+    c = { :host => "host" }
+    @session.expects(:open_channel).yields(c).returns(c)
+    c.expects(:exec).with('something').yields(c, true)
+    c.expects(:on_data).yields(c, "stdout")
+    c.expects(:on_extended_data).yields(c, 1, "stderr")
+    c.expects(:on_request)
+    results = {}
+    @session.exec("something") do |c, stream, data|
+      results[stream] = data
+    end
+    assert_equal({:stdout => "stdout", :stderr => "stderr"}, results)
+  end
+
+  def test_exec_without_block_should_write_data_and_extended_data_lines_to_stdout_and_stderr
+    c = { :host => "host" }
+    @session.expects(:open_channel).yields(c).returns(c)
+    c.expects(:exec).with('something').yields(c, true)
+    c.expects(:on_data).yields(c, "stdout 1\nstdout 2\n")
+    c.expects(:on_extended_data).yields(c, 1, "stderr 1\nstderr 2\n")
+    c.expects(:on_request)
+    $stdout.expects(:puts).with("[host] stdout 1\n")
+    $stdout.expects(:puts).with("[host] stdout 2")
+    $stderr.expects(:puts).with("[host] stderr 1\n")
+    $stderr.expects(:puts).with("[host] stderr 2")
+    @session.exec("something")
+  end
+
+  def test_exec_should_capture_exit_status_of_process
+    c = { :host => "host" }
+    @session.expects(:open_channel).yields(c).returns(c)
+    c.expects(:exec).with('something').yields(c, true)
+    c.expects(:on_data)
+    c.expects(:on_extended_data)
+    c.expects(:on_request).with("exit-status").yields(c, Net::SSH::Buffer.from(:long, 127))
+    @session.exec("something")
+    assert_equal 127, c[:exit_status]
+  end
+
+end
\ No newline at end of file
diff --git a/test/session_test.rb b/test/session_test.rb
index 8aebccc..288275a 100644
--- a/test/session_test.rb
+++ b/test/session_test.rb
@@ -73,85 +73,66 @@ class SessionTest < Test::Unit::TestCase
     assert_equal s1.object_id, s2.object_id
   end
 
-  def test_with_should_set_active_groups_and_yield_and_restore_active_groups
+  def test_with_should_yield_new_subsession_with_servers_for_criteria
     yielded = nil
-    @session.with(:app, :web) do |s|
+    @session.expects(:servers_for).with(:app, :web).returns([:servers])
+    result = @session.with(:app, :web) do |s|
       yielded = s
-      assert_equal({:app => {}, :web => {}}, @session.active_groups)
     end
-    assert_equal @session, yielded
-    assert_equal({}, @session.active_groups)
+    assert_equal result, yielded
+    assert_equal [:servers], yielded.servers
   end
 
-  def test_with_with_unknown_constraint_should_raise_error
+  def test_servers_for_with_unknown_constraint_should_raise_error
     assert_raises(ArgumentError) do
-      @session.with(:app => { :all => :foo }) {}
+      @session.servers_for(:app => { :all => :foo })
     end
   end
 
-  def test_with_with_constraints_should_add_constraints_to_active_groups
-    @session.with(:app => { :only => { :primary => true }, :except => { :backup => true } }) do |s|
-      assert_equal({:app => {:only => {:primary => true}, :except => {:backup => true}}}, @session.active_groups)
-    end
+  def test_with_with_constraints_should_build_subsession_with_matching_servers
+    conditions = { :app => { :only => { :primary => true }, :except => { :backup => true } } }
+    @session.expects(:servers_for).with(conditions).returns([:servers])
+    assert_equal [:servers], @session.with(conditions).servers
   end
 
-  def test_on_should_create_ad_hoc_group_and_make_that_group_the_only_active_group
+  def test_on_should_return_subsession_containing_only_the_given_servers
     s1 = @session.use('h1', 'u1')
     s2 = @session.use('h2', 'u2')
+    subsession = @session.on(s1, s2)
+    assert_equal [s1, s2], subsession.servers
+  end
+
+  def test_on_should_yield_subsession_if_block_is_given
+    s1 = @session.use('h1', 'u1')
     yielded = nil
-    @session.active_groups[:g1] = []
-    @session.on(s1, s2) do |s|
+    result = @session.on(s1) do |s|
       yielded = s
-      assert_equal 1, @session.active_groups.size
-      assert_not_equal :g1, @session.active_groups.keys.first
-      assert_equal [s1, s2], @session.groups[@session.active_groups.keys.first]
+      assert_equal [s1], s.servers
     end
-    assert_equal [:g1], @session.active_groups.keys
-    assert_equal @session, yielded
+    assert_equal result, yielded
   end
 
-  def test_active_sessions_should_return_sessions_for_all_servers_if_active_groups_is_empty
-    s1, s2, s3 = MockSession.new, MockSession.new, MockSession.new
+  def test_servers_for_should_return_all_servers_if_no_arguments
     srv1, srv2, srv3 = @session.use('h1', 'u1'), @session.use('h2', 'u2'), @session.use('h3', 'u3')
-    Net::SSH.expects(:start).with('h1', 'u1', {}).returns(s1)
-    Net::SSH.expects(:start).with('h2', 'u2', {}).returns(s2)
-    Net::SSH.expects(:start).with('h3', 'u3', {}).returns(s3)
-    assert_equal [s1, s2, s3], @session.active_sessions.sort
+    assert_equal %w(h1 h2 h3), @session.servers_for.map { |s| s.host }.sort
   end
 
-  def test_active_sessions_should_return_sessions_only_for_active_groups_if_active_groups_exist
-    s1, s2, s3 = MockSession.new, MockSession.new, MockSession.new
+  def test_servers_for_should_return_servers_only_for_given_group
     srv1, srv2, srv3 = @session.use('h1', 'u1'), @session.use('h2', 'u2'), @session.use('h3', 'u3')
     @session.group :app => [srv1, srv2], :db => [srv3]
-    Net::SSH.expects(:start).with('h1', 'u1', {}).returns(s1)
-    Net::SSH.expects(:start).with('h2', 'u2', {}).returns(s2)
-    @session.active_groups.replace(:app => {})
-    assert_equal [s1, s2], @session.active_sessions.sort
+    assert_equal %w(h1 h2), @session.servers_for(:app).map { |s| s.host }.sort
   end
 
-  def test_active_sessions_should_not_return_duplicate_sessions
-    s1, s2, s3 = MockSession.new, MockSession.new, MockSession.new
+  def test_servers_for_should_not_return_duplicate_servers
     srv1, srv2, srv3 = @session.use('h1', 'u1'), @session.use('h2', 'u2'), @session.use('h3', 'u3')
     @session.group :app => [srv1, srv2], :db => [srv2, srv3]
-    Net::SSH.expects(:start).with('h1', 'u1', {}).returns(s1)
-    Net::SSH.expects(:start).with('h2', 'u2', {}).returns(s2)
-    Net::SSH.expects(:start).with('h3', 'u3', {}).returns(s3)
-    @session.active_groups.replace(:app => {}, :db => {})
-    assert_equal [s1, s2, s3], @session.active_sessions.sort
+    assert_equal ["h1", "h2", "h3"], @session.servers_for(:app, :db).map { |s| s.host }.sort
   end
 
-  def test_active_sessions_should_correctly_apply_only_and_except_constraints
-    s1, s2, s3 = MockSession.new, MockSession.new, MockSession.new
+  def test_servers_for_should_correctly_apply_only_and_except_constraints
     srv1, srv2, srv3 = @session.use('h1', 'u1', :properties => {:a => 1}), @session.use('h2', 'u2', :properties => {:a => 1, :b => 2}), @session.use('h3', 'u3')
     @session.group :app => [srv1, srv2, srv3]
-    Net::SSH.expects(:start).with('h1', 'u1', :properties => {:a => 1}).returns(s1)
-    @session.active_groups.replace(:app => {:only => {:a => 1}, :except => {:b => 2}})
-    assert_equal [s1], @session.active_sessions.sort
-  end
-
-  def test_connect_bang_should_call_active_sessions_and_return_self
-    @session.expects(:active_sessions)
-    assert_equal @session, @session.connect!
+    assert_equal [srv1], @session.servers_for(:app => {:only => {:a => 1}, :except => {:b => 2}})
   end
 
   def test_close_should_close_server_sessions
@@ -171,22 +152,6 @@ class SessionTest < Test::Unit::TestCase
     @session.close
   end
 
-  def test_busy_should_be_true_if_any_server_is_busy
-    srv1, srv2, srv3 = @session.use('h1', 'u1', :properties => {:a => 1}), @session.use('h2', 'u2', :properties => {:a => 1, :b => 2}), @session.use('h3', 'u3')
-    srv1.stubs(:busy?).returns(false)
-    srv2.stubs(:busy?).returns(false)
-    srv3.stubs(:busy?).returns(true)
-    assert @session.busy?
-  end
-
-  def test_busy_should_be_false_if_all_servers_are_not_busy
-    srv1, srv2, srv3 = @session.use('h1', 'u1', :properties => {:a => 1}), @session.use('h2', 'u2', :properties => {:a => 1, :b => 2}), @session.use('h3', 'u3')
-    srv1.stubs(:busy?).returns(false)
-    srv2.stubs(:busy?).returns(false)
-    srv3.stubs(:busy?).returns(false)
-    assert !@session.busy?
-  end
-
   def test_loop_should_loop_until_process_is_false
     @session.expects(:process).with(5).times(4).returns(true,true,true,false).yields
     yielded = false
@@ -234,103 +199,4 @@ class SessionTest < Test::Unit::TestCase
     IO.expects(:select).with([:a, :b, :c], [:a, :c], nil, 5).returns([[:b, :c], [:a, :c]])
     @session.process(5)
   end
-
-  def test_send_global_request_should_delegate_to_active_sessions
-    s1 = mock('ssh')
-    s2 = mock('ssh')
-    s1.expects(:send_global_request).with("a", "b", "c").yields
-    s2.expects(:send_global_request).with("a", "b", "c").yields
-    @session.expects(:active_sessions).returns([s1, s2])
-    calls = 0
-    @session.send_global_request("a", "b", "c") { calls += 1 }
-    assert_equal 2, calls
-  end
-
-  def test_open_channel_should_delegate_to_active_sessions_and_set_accessors_on_each_channel_and_return_multi_channel
-    srv1 = @session.use('h1', 'u1')
-    srv2 = @session.use('h2', 'u2')
-    s1 = { :server => srv1 }
-    s2 = { :server => srv2 }
-    c1 = { :stub => :value }
-    c2 = {}
-    c1.stubs(:connection).returns(s1)
-    c2.stubs(:connection).returns(s2)
-    @session.expects(:active_sessions).returns([s1, s2])
-    s1.expects(:open_channel).with("session").yields(c1).returns(c1)
-    s2.expects(:open_channel).with("session").yields(c2).returns(c2)
-    results = []
-    channel = @session.open_channel do |c|
-      results << c
-    end
-    assert_equal [c1, c2], results
-    assert_equal "h1", c1[:host]
-    assert_equal "h2", c2[:host]
-    assert_equal srv1, c1[:server]
-    assert_equal srv2, c2[:server]
-    assert_instance_of Net::SSH::Multi::Channel, channel
-    assert_equal [c1, c2], channel.channels
-  end
-
-  def test_exec_should_raise_exception_if_channel_cannot_exec_command
-    c = { :host => "host" }
-    @session.expects(:open_channel).yields(c).returns(c)
-    c.expects(:exec).with('something').yields(c, false)
-    assert_raises(RuntimeError) { @session.exec("something") }
-  end
-
-  def test_exec_with_block_should_pass_data_and_extended_data_to_block
-    c = { :host => "host" }
-    @session.expects(:open_channel).yields(c).returns(c)
-    c.expects(:exec).with('something').yields(c, true)
-    c.expects(:on_data).yields(c, "stdout")
-    c.expects(:on_extended_data).yields(c, 1, "stderr")
-    c.expects(:on_request)
-    results = {}
-    @session.exec("something") do |c, stream, data|
-      results[stream] = data
-    end
-    assert_equal({:stdout => "stdout", :stderr => "stderr"}, results)
-  end
-
-  def test_exec_without_block_should_write_data_and_extended_data_lines_to_stdout_and_stderr
-    c = { :host => "host" }
-    @session.expects(:open_channel).yields(c).returns(c)
-    c.expects(:exec).with('something').yields(c, true)
-    c.expects(:on_data).yields(c, "stdout 1\nstdout 2\n")
-    c.expects(:on_extended_data).yields(c, 1, "stderr 1\nstderr 2\n")
-    c.expects(:on_request)
-    $stdout.expects(:puts).with("[host] stdout 1\n")
-    $stdout.expects(:puts).with("[host] stdout 2")
-    $stderr.expects(:puts).with("[host] stderr 1\n")
-    $stderr.expects(:puts).with("[host] stderr 2")
-    @session.exec("something")
-  end
-
-  def test_exec_should_capture_exit_status_of_process
-    c = { :host => "host" }
-    @session.expects(:open_channel).yields(c).returns(c)
-    c.expects(:exec).with('something').yields(c, true)
-    c.expects(:on_data)
-    c.expects(:on_extended_data)
-    c.expects(:on_request).with("exit-status").yields(c, Net::SSH::Buffer.from(:long, 127))
-    @session.exec("something")
-    assert_equal 127, c[:exit_status]
-  end
-
-  private
-
-    class MockSession < Hash
-      include Comparable
-
-      @@next_id = 0
-      attr_reader :id
-
-      def initialize
-        @id = (@@next_id += 1)
-      end
-
-      def <=>(s)
-        id <=> s.id
-      end
-    end
 end
\ No newline at end of file