session.rb in net-ssh-session-0.1.5

- old
+ new

@@ -12,10 +12,11 @@
       attr_reader :host, :user, :password
       attr_reader :connection, :shell
       attr_reader :options
       attr_reader :logger
       attr_reader :history
+      attr_reader :timeout
 
       # Initialize a new ssh session
       # @param [String] remote hostname or ip address
       # @param [String] remote account username
       # @param [String] remote account password
@@ -24,14 +25,19 @@
         @host          = host
         @user          = user
         @password      = password
         @history       = []
         @track_history = true
+        @timeout       = options[:timeout]
 
         if options[:history] == false
           @track_history = false
         end
+
+        if @timeout && !@timeout.kind_of?(Integer)
+          raise ArgumentError, "Timeout value should be numeric"
+        end
       end
 
       # Establish connection with remote server
       # @param [Integer] max timeout in seconds
       # @return [Boolean]
@@ -57,16 +63,21 @@
       # @param [String] command to execute
       # @param [Block] output event block
       # @return [Integer] command execution exit code
       def exec(command, &on_output)
         status = nil
-        shell.execute(command) do |process|
-          process.on_output(&on_output)
-          process.on_error_output(&on_output)
-          process.on_finish { |p| status = p.exit_status }
+
+        handle_timeout do
+          shell.execute(command) do |process|
+            process.on_output(&on_output)
+            process.on_error_output(&on_output)
+            process.on_finish { |p| status = p.exit_status }
+          end
+
+          shell.session.loop(1) { status.nil? }
         end
-        shell.session.loop(1) { status.nil? }
+
         status
       end
 
       # Execute a single command
       # @param [String] comand to execute
@@ -131,18 +142,29 @@
       # @return [SessionCommand]
       def last_command
         history.last
       end
 
+      # Execute a dynamic command
+      # @param [String] command name
+      # @params [Array] command arguments
       def method_missing(name, *args)
         run("#{name} #{args.join(' ')}".strip)
       end
 
       private
 
       def establish_connection
         @connection = Net::SSH.start(host, user, :password => password)
         @shell = @connection.shell
+      end
+
+      def handle_timeout(&block)
+        if timeout
+          Timeout.timeout(timeout) { block.call(self) }
+        else
+          block.call(self)
+        end
       end
     end
   end
 end
\ No newline at end of file