From 9782ac60ccc82169a0e942c7d0eaba008bae44e8 Mon Sep 17 00:00:00 2001 From: Jamis Buck Date: Mon, 7 Apr 2008 22:41:12 -0600 Subject: make with() and on() yield a new subsession object that encapsulates the set of matching servers, and add a new servers_for method. --- README.rdoc | 4 +- lib/net/ssh/multi.rb | 4 +- lib/net/ssh/multi/session.rb | 285 +++++++++-------------------------- lib/net/ssh/multi/session_actions.rb | 153 +++++++++++++++++++ lib/net/ssh/multi/subsession.rb | 48 ++++++ test/session_actions_test.rb | 124 +++++++++++++++ test/session_test.rb | 192 ++++------------------- 7 files changed, 430 insertions(+), 380 deletions(-) create mode 100644 lib/net/ssh/multi/session_actions.rb create mode 100644 lib/net/ssh/multi/subsession.rb create mode 100644 test/session_actions_test.rb 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 -- cgit v1.2.1