Merge branch 'master' of github.com:fauna/thrift_client

lib/thrift_client/abstract_thrift_client.rb

@@ -52,6 +52,7 @@ def initialize(client_class, servers, options = {})
     @server_list = Array(servers).collect{|s| Server.new(s)}.sort_by { rand }
     @current_server = @server_list.first
 
+    @callbacks = {}
     @client_methods = []
     @client_class.instance_methods.each do |method_name|
       if method_name != 'send_message' && method_name =~ /^send_(.*)$/
@@ -70,6 +71,16 @@ def initialize(client_class, servers, options = {})
     end
   end
 
+  def add_callback(cb, &blk)
+    if cb == :post_connect
+      @callbacks[cb] = blk
+      # Allow chaining
+      return self
+    else
+      return nil
+    end
+  end
+
   def inspect
     "<#{self.class}(#{client_class}) @current_server=#{@current_server}>"
   end
@@ -84,6 +95,7 @@ def connect!
     transport = @connection.transport
     transport.timeout = @options[:timeout] if transport_can_timeout?
     @client = @client_class.new(@options[:protocol].new(transport, *@options[:protocol_extra_params]))
+    @callbacks[:post_connect].call(self) if @callbacks[:post_connect]
   end
 
   def disconnect!

test/thrift_client_test.rb

@@ -83,12 +83,42 @@ def test_lazy_connection
     end
   end
 
+  def test_post_conn_cb
+    calledcnt = 0
+    client = ThriftClient.new(Greeter::Client, @servers, @options.merge(:retries => 2))
+    r = client.add_callback :post_connect do |cl|
+      calledcnt += 1
+      assert_equal(client, cl)
+    end
+    assert_equal(client, r)
+    assert_nothing_raised do
+      client.greeting("someone")
+      client.disconnect!
+    end
+    assert_equal(1, calledcnt)
+  end
+
+
+  def test_unknown_cb
+    calledcnt = 0
+    client = ThriftClient.new(Greeter::Client, @servers, @options.merge(:retries => 2))
+    r = client.add_callback :unknown do |cl|
+      assert(false)
+    end
+    assert_equal(nil, r)
+  end
+
   def test_no_servers_eventually_raise
+    wascalled = false
     client = ThriftClient.new(Greeter::Client, @servers[0,2], @options.merge(:retries => 2))
+    client.add_callback :post_connect do
+      wascalled = true
+    end
     assert_raises(ThriftClient::NoServersAvailable) do
       client.greeting("someone")
       client.disconnect!
     end
+    assert(!wascalled)
   end
 
   def test_socket_timeout