lib/dripdrop/handlers/http.rb in dripdrop-0.9.3 vs lib/dripdrop/handlers/http.rb in dripdrop-0.9.4

@@ -15,11 +15,11 @@
       @body_callback = blk
     end
     
     def send_message(raw_msg)
       msg = dd_messagify(raw_msg)
-      if msg.class == DripDrop::Message
+      if msg.is_a?(DripDrop::Message)
         json = msg.encode_json
         self.call([json])
         self.succeed
       else
         raise "Message Type '#{msg.class}' not supported"
@@ -61,10 +61,11 @@
     
     def initialize(uri,opts={})
       @uri     = uri
       @address = uri.to_s
       @opts    = opts
+      @message_class = @opts[:message_class] || DripDrop.default_message_class
     end
     
     def on_recv(msg_format=:dripdrop_json,&block)
       #Thin's error handling only rescues stuff w/o a backtrace
       begin
@@ -88,24 +89,25 @@
     
     def initialize(uri, opts={})
       @uri     = uri
       @address = @uri.to_s
       @opts    = opts
+      @message_class = @opts[:message_class] || DripDrop.default_message_class
     end
     
     def send_message(message,&block)
       dd_message = dd_messagify(message)
-      if dd_message.class == DripDrop::Message
+      if dd_message.is_a?(DripDrop::Message)
         uri_path = @uri.path.empty? ? '/' : @uri.path
         
         req = EM::Protocols::HttpClient.request(
           :host => @uri.host, :port => @uri.port,
           :request => uri_path, :verb => 'POST',
           :contenttype => 'application/json',
           :content => dd_message.encode_json
         )
         req.callback do |response|
-          block.call(DripDrop::Message.decode_json(response[:content]))
+          block.call(@message_class.decode_json(response[:content]))
         end
       else
         raise "Unsupported message type '#{dd_message.class}'"
       end
     end