✨ Add response_handlers option to new (backport)

lib/net/imap.rb

@@ -134,6 +134,8 @@ module Net
   # Use paginated or limited versions of commands whenever possible.
   #
   # Use #add_response_handler to handle responses after each one is received.
+  # Use the +response_handlers+ argument to ::new to assign response handlers
+  # before the receiver thread is started.
   #
   # == Errors
   #
@@ -984,6 +986,11 @@ def uid_sort(sort_keys, search_keys, charset)
     #     end
     #   }
     #
+    # Response handlers can also be added when the client is created before the
+    # receiver thread is started, by the +response_handlers+ argument to ::new.
+    # This ensures every server response is handled, including the #greeting.
+    #
+    # Related: #remove_response_handler, #response_handlers
     def add_response_handler(handler = nil, &block)
       raise ArgumentError, "two Procs are passed" if handler && block
       @response_handlers.push(block || handler)
@@ -1099,6 +1106,12 @@ def idle_done
     #         OpenSSL::SSL::SSLContext#set_params as parameters.
     # open_timeout:: Seconds to wait until a connection is opened
     # idle_response_timeout:: Seconds to wait until an IDLE response is received
+    # response_handlers:: A list of response handlers to be added before the
+    #                     receiver thread is started.  This ensures every server
+    #                     response is handled, including the #greeting.  Note
+    #                     that the greeting is handled in the current thread,
+    #                     but all other responses are handled in the receiver
+    #                     thread.
     #
     # The most common errors are:
     #
@@ -1141,6 +1154,7 @@ def initialize(host, port_or_options = {},
         @responses = Hash.new([].freeze)
         @tagged_responses = {}
         @response_handlers = []
+        options[:response_handlers]&.each do |h| add_response_handler(h) end
         @tagged_response_arrival = new_cond
         @continued_command_tag = nil
         @continuation_request_arrival = new_cond
@@ -1157,6 +1171,7 @@ def initialize(host, port_or_options = {},
         if @greeting.name == "BYE"
           raise ByeResponseError, @greeting
         end
+        @response_handlers.each do |handler| handler.call(@greeting) end
 
         @client_thread = Thread.current
         @receiver_thread = Thread.start {

test/net/imap/test_imap_response_handlers.rb

@@ -0,0 +1,134 @@
+# frozen_string_literal: true
+
+require "net/imap"
+require "test/unit"
+
+class IMAPResponseHandlersTest < Test::Unit::TestCase
+
+  def setup
+    @do_not_reverse_lookup = Socket.do_not_reverse_lookup
+    Socket.do_not_reverse_lookup = true
+    @threads = []
+  end
+
+  def teardown
+    if [email protected]?
+      assert_join_threads(@threads)
+    end
+  ensure
+    Socket.do_not_reverse_lookup = @do_not_reverse_lookup
+  end
+
+  test "#add_response_handlers" do
+    server = create_tcp_server
+    port   = server.addr[1]
+    start_server do
+      sock = server.accept
+      Timeout.timeout(5) do
+        sock.print("* OK connection established\r\n")
+        sock.gets # => NOOP
+        sock.print("* 1 EXPUNGE\r\n")
+        sock.print("* 2 EXPUNGE\r\n")
+        sock.print("* 3 EXPUNGE\r\n")
+        sock.print("RUBY0001 OK NOOP completed\r\n")
+        sock.gets # => LOGOUT
+        sock.print("* BYE terminating connection\r\n")
+        sock.print("RUBY0002 OK LOGOUT completed\r\n")
+      ensure
+        sock.close
+        server.close
+      end
+    end
+    begin
+      responses = []
+      imap = Net::IMAP.new(server_addr, port: port)
+      assert_equal 0, imap.response_handlers.length
+      imap.add_response_handler do |r| responses << [:block, r] end
+      assert_equal 1, imap.response_handlers.length
+      imap.add_response_handler(->(r) { responses << [:proc, r] })
+      assert_equal 2, imap.response_handlers.length
+
+      imap.noop
+      responses = responses[0, 6].map {|which, resp|
+        [which, resp.class, resp.name, resp.data]
+      }
+      assert_equal [
+        [:block, Net::IMAP::UntaggedResponse, "EXPUNGE", 1],
+        [:proc,  Net::IMAP::UntaggedResponse, "EXPUNGE", 1],
+        [:block, Net::IMAP::UntaggedResponse, "EXPUNGE", 2],
+        [:proc,  Net::IMAP::UntaggedResponse, "EXPUNGE", 2],
+        [:block, Net::IMAP::UntaggedResponse, "EXPUNGE", 3],
+        [:proc,  Net::IMAP::UntaggedResponse, "EXPUNGE", 3],
+      ], responses
+    ensure
+      imap&.logout
+      imap&.disconnect
+    end
+  end
+
+  test "::new with response_handlers kwarg" do
+    greeting = nil
+    expunges = []
+    alerts   = []
+    untagged = 0
+    handler0 = ->(r) { greeting ||= r }
+    handler1 = ->(r) { alerts   << r.data.text if r.data.code.name == "ALERT" rescue nil }
+    handler2 = ->(r) { expunges << r.data if r.name == "EXPUNGE" }
+    handler3 = ->(r) { untagged += 1 if r.is_a?(Net::IMAP::UntaggedResponse) }
+    response_handlers = [handler0, handler1, handler2, handler3]
+
+    server = create_tcp_server
+    port   = server.addr[1]
+    start_server do
+      sock = server.accept
+      Timeout.timeout(5) do
+        sock.print("* OK connection established\r\n")
+        sock.gets # => NOOP
+        sock.print("* 1 EXPUNGE\r\n")
+        sock.print("* 1 EXPUNGE\r\n")
+        sock.print("* OK [ALERT] The first alert.\r\n")
+        sock.print("RUBY0001 OK [ALERT] Did you see the alert?\r\n")
+        sock.gets # => LOGOUT
+        sock.print("* BYE terminating connection\r\n")
+        sock.print("RUBY0002 OK LOGOUT completed\r\n")
+      ensure
+        sock.close
+        server.close
+      end
+    end
+    begin
+      imap = Net::IMAP.new("localhost", port: port,
+                           response_handlers: response_handlers)
+      assert_equal response_handlers, imap.response_handlers
+      refute_same  response_handlers, imap.response_handlers
+
+      # handler0 recieved the greeting and handler3 counted it
+      assert_equal imap.greeting, greeting
+      assert_equal 1, untagged
+
+      imap.noop
+      assert_equal 4, untagged
+      assert_equal [1, 1], expunges # from handler2
+      assert_equal ["The first alert.", "Did you see the alert?"], alerts
+    ensure
+      imap&.logout
+      imap&.disconnect
+    end
+  end
+
+  def start_server
+    th = Thread.new do
+      yield
+    end
+    @threads << th
+    sleep 0.1 until th.stop?
+  end
+
+  def create_tcp_server
+    return TCPServer.new(server_addr, 0)
+  end
+
+  def server_addr
+    Addrinfo.tcp("localhost", 0).ip_address
+  end
+end