Merge pull request #225 from mfazekas/reafactor-proxy-code

Refactor proxy code

2 files changed, 9 insertions, 8 deletions

diff --git a/lib/net/ssh/transport/session.rb b/lib/net/ssh/transport/session.rb
index 76dd18b..19e6bbd 100644
--- a/lib/net/ssh/transport/session.rb
+++ b/lib/net/ssh/transport/session.rb
@@ -63,14 +63,15 @@ module Net; module SSH; module Transport
       @options = options
 
       debug { "establishing connection to #{@host}:#{@port}" }
-      factory = options[:proxy] || TCPSocket
-      @socket = timeout(options[:timeout] || 0) {
-        case
-        when options[:proxy] then factory.open(@host, @port, options)
-        when @bind_address.nil? then factory.open(@host, @port)
-        else factory.open(@host, @port, @bind_address)
+
+      @socket = timeout(options[:timeout] || 0) do
+        if (factory = options[:proxy])
+          factory.open(@host, @port, options)
+        else
+          TCPSocket.open(@host, @port, @bind_address)
         end
-      }
+      end
+
       @socket.extend(PacketStream)
       @socket.logger = @logger
 
diff --git a/test/transport/test_session.rb b/test/transport/test_session.rb
index 88e8402..6c267ec 100644
--- a/test/transport/test_session.rb
+++ b/test/transport/test_session.rb
@@ -315,7 +315,7 @@ module Transport
       def session(options={})
         @session ||= begin
           host = options.delete(:host) || "net.ssh.test"
-          TCPSocket.stubs(:open).with(host, options[:port] || 22).returns(socket)
+          TCPSocket.stubs(:open).with(host, options[:port] || 22, nil).returns(socket)
           Net::SSH::Transport::ServerVersion.stubs(:new).returns(server_version)
           Net::SSH::Transport::Algorithms.stubs(:new).returns(algorithms)