server selection based on properties as well as groups

2 files changed, 31 insertions, 15 deletions

diff --git a/lib/net/ssh/multi/server.rb b/lib/net/ssh/multi/server.rb
index c3eae2e..252304b 100644
--- a/lib/net/ssh/multi/server.rb
+++ b/lib/net/ssh/multi/server.rb
@@ -14,6 +14,10 @@ module Net; module SSH; module Multi
       @gateway = @options.delete(:via)
     end
 
+    def [](key)
+      (@options[:properties] || {})[key]
+    end
+
     def port
       options[:port] || 22
     end
diff --git a/lib/net/ssh/multi/session.rb b/lib/net/ssh/multi/session.rb
index efcd7fe..9d44e71 100644
--- a/lib/net/ssh/multi/session.rb
+++ b/lib/net/ssh/multi/session.rb
@@ -15,7 +15,8 @@ module Net; module SSH; module Multi
       @gateway = nil
       @connections_mutex = Mutex.new
       @groups_mutex = Mutex.new
-      @active_groups = []
+      @active_groups = {}
+      @open_groups = []
     end
 
     def group(*args)
@@ -25,15 +26,15 @@ module Net; module SSH; module Multi
         raise ArgumentError, "must provide group mapping OR block, not both"
       elsif block_given?
         begin
-          saved_groups = active_groups.dup
-          active_groups.concat(args.map { |a| a.to_sym }).uniq!
+          saved_groups = @open_groups.dup
+          @open_groups.concat(args.map { |a| a.to_sym }).uniq!
           yield self if block_given?
         ensure
-          active_groups.replace(saved_groups)
+          @open_groups.replace(saved_groups)
         end
       else
         mapping.each do |key, value|
-          (active_groups + Array(key)).uniq.each do |grp|
+          (@open_groups + Array(key)).uniq.each do |grp|
             (groups[grp.to_sym] ||= []).concat(Array(value))
           end
         end
@@ -55,21 +56,36 @@ module Net; module SSH; module Multi
     end
 
     def with(*groups)
-      saved_groups = active_groups.dup
-      active_groups.concat(groups).uniq!
+      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)
+            map[gr] = (@active_groups[gr] || {}).merge(value)
+          end
+        else
+          map[group] = @active_groups[group] || {}
+        end
+        map
+      end
+
+      @active_groups.update(new_map)
       yield self
     ensure
-      active_groups.replace(saved_groups)
+      @active_groups.replace(saved_groups)
     end
 
     def active_sessions
-      list = if active_groups.empty?
+      list = if @active_groups.empty?
         servers
       else
-        active_groups.map { |group| groups[group] }.flatten.uniq
+        @active_groups.inject([]) do |list, (group, properties)|
+          list.concat(groups[group].select { |server| properties.all? { |prop, value| server[prop] == value } })
+        end
       end
 
-      sessions_for(list)
+      sessions_for(list.uniq)
     end
 
     def connect!
@@ -168,10 +184,6 @@ module Net; module SSH; module Multi
 
     private
 
-      def active_groups
-        @active_groups
-      end
-
       def sessions_for(servers)
         threads = servers.map { |server| Thread.new { server.session(true) } }
         threads.each { |thread| thread.join }