test_client.rb in ftpd-0.6.0

- old
+ new

@@ -4,17 +4,24 @@
 class TestClient
 
   extend Forwardable
   include FileUtils
 
-  def initialize(opts = {})
+  attr_accessor :tls_mode
+
+  def initialize
+    @tls_mode = :off
     @temp_dir = Ftpd::TempDir.make
-    @ftp = make_ftp(opts)
     @templates = TestFileTemplates.new
   end
 
+  def start
+    @ftp = make_ftp
+  end
+
   def close
+    return unless @ftp
     @ftp.close
   end
 
   def_delegators :@ftp,
   :chdir,
@@ -34,10 +41,24 @@
   :rename,
   :rmdir,
   :status,
   :system
 
+  # Make a connection from a specific IP.  Net::FTP doesn't have a way
+  # to force the local IP, so fake it here.
+
+  def connect_from(source_ip, host, port)
+    in_addr = Socket.pack_sockaddr_in(0, source_ip)
+    out_addr = Socket.pack_sockaddr_in(port, host)
+    socket = Socket.open(Socket::AF_INET, Socket::SOCK_STREAM, 0)
+    socket.bind(in_addr)
+    socket.connect(out_addr)
+    decorate_socket socket
+    @ftp = make_ftp
+    @ftp.set_socket(socket)
+  end
+
   def raw(*command)
     @ftp.sendcmd command.compact.join(' ')
   end
 
   def get(mode, remote_path)
@@ -65,11 +86,11 @@
   def file_contents(path)
     File.open(temp_path(path), 'rb', &:read)
   end
 
   def xpwd
-    response = raw('XPWD')
+    response = @ftp.sendcmd('XPWD')
     response[/"(.+)"/, 1]
   end
 
   def store_unique(local_path, remote_path)
     command = ['STOU', remote_path].compact.join(' ')
@@ -101,10 +122,14 @@
     rescue EOFError
       false
     end
   end
 
+  def set_option(option)
+    @ftp.sendcmd "OPTS #{option}"
+  end
+
   private
 
   RAW_METHOD_REGEX = /^send_(.*)$/
 
   def local_path(remote_path)
@@ -113,21 +138,20 @@
 
   def temp_path(path)
     File.expand_path(path, @temp_dir)
   end
 
-  def make_ftp(opts)
-    tls_mode = opts[:tls] || :off
-    case tls_mode
+  def make_ftp
+    case @tls_mode
     when :off
       make_non_tls_ftp
     when :implicit
       make_tls_ftp(:implicit)
     when :explicit
       make_tls_ftp(:explicit)
     else
-      raise "Unknown TLS mode: #{tls_mode}"
+      raise "Unknown TLS mode: #{@tls_mode}"
     end
   end
 
   def make_tls_ftp(ftps_mode)
     ftp = DoubleBagFTPS.new
@@ -139,8 +163,39 @@
     ftp
   end
 
   def make_non_tls_ftp
     Net::FTP.new
+  end
+
+  # Ruby 2.0's Ftp class is expecting a TCPSocket, not a Socket.  The
+  # trouble comes with Ftp#close, which closes sockets by first doing
+  # a shutdown, setting the read timeout, and doing a read.  Plain
+  # Socket doesn't have those methods, so fake it.
+  #
+  # Plain socket _does_ have #close, but we short-circuit it, too,
+  # because it takes a few seconds.  We're in a hurry when running
+  # tests, and can afford to be a little sloppy when cleaning up.
+
+  def decorate_socket(sock)
+
+    def sock.shutdown(how)
+      @shutdown = true
+    end
+
+    def sock.read_timeout=(seconds)
+    end
+
+    # Skip read after shutdown.  Prevents 2.0 from hanging in
+    # Ftp#close
+
+    def sock.read(*args)
+      return if @shutdown
+      super(*args)
+    end
+
+    def close
+    end
+
   end
 
 end