Merge pull request #2 from stekker/connection-callbacks

Implement Connection callbacks

lib/s2-ruby.rb

@@ -10,6 +10,7 @@
 require_relative "s2/messages/reception_status"
 require_relative "s2/message_factory"
 require_relative "s2/message_handler"
+require_relative "s2/message_handler_callbacks"
 require_relative "s2/connection"
 
 module S2

lib/s2/connection.rb

@@ -1,6 +1,7 @@
 module S2
   class Connection
     include S2::MessageHandler
+    include S2::MessageHandlerCallbacks
 
     attr_reader :state, :sent_messages
 
@@ -17,11 +18,18 @@ def initialize(ws, logger: Rails.logger)
       @ws = ws
       @logger = logger
       @sent_messages = {}
+      @rm_id = nil
 
       update_state :connected
     end
 
+    def open(rm_id)
+      @rm_id = rm_id
+      trigger_on_open(rm_id)
+    end
+
     def receive_message(message_json)
+      trigger_before_receive(@rm_id, message_json)
       @logger.info("Received message: #{message_json}")
 
       message = deserialize_message(message_json)
@@ -37,6 +45,7 @@ def send_message(message_class, payload)
       store_sent_message(message) unless message.is_a?(S2::Messages::ReceptionStatus)
       json = message.to_json
       send_raw_message(json)
+      trigger_after_send(@rm_id, json)
 
       @logger.info("Sent message: #{json}")
     end

lib/s2/message_handler_callbacks.rb

@@ -0,0 +1,39 @@
+module S2
+  module MessageHandlerCallbacks
+    extend ActiveSupport::Concern
+
+    included do
+      class_attribute :on_open_proc, default: nil
+      class_attribute :before_receive_proc, default: nil
+      class_attribute :after_send_proc, default: nil
+    end
+
+    class_methods do
+      def on_open(&block)
+        self.on_open_proc = block
+      end
+
+      def before_receive(&block)
+        self.before_receive_proc = block
+      end
+
+      def after_send(&block)
+        self.after_send_proc = block
+      end
+    end
+
+    protected
+
+    def trigger_on_open(rm_id)
+      instance_exec(rm_id, &self.class.on_open_proc) if self.class.on_open_proc
+    end
+
+    def trigger_before_receive(rm_id, message)
+      instance_exec(rm_id, message, &self.class.before_receive_proc) if self.class.before_receive_proc
+    end
+
+    def trigger_after_send(rm_id, message)
+      instance_exec(rm_id, message, &self.class.after_send_proc) if self.class.after_send_proc
+    end
+  end
+end

spec/lib/s2/connection_spec.rb

@@ -102,4 +102,62 @@
       expect(ws).to be_closed
     end
   end
+
+  describe "callbacks" do
+    it "calls a callback on opening the connection" do
+      connection_class = Class.new(described_class) do
+        attr_reader :on_open_arg
+
+        on_open do |rm_id|
+          @on_open_arg = rm_id
+        end
+      end
+
+      connection = connection_class.new(ws, logger: Logger.new(nil))
+
+      expect { connection.open("123") }
+        .to change { connection.on_open_arg }.from(nil).to("123")
+    end
+
+    it "calls a callback before handling a message" do
+      connection_class = Class.new(described_class) do
+        attr_reader :before_receive_args
+
+        before_receive do |*args|
+          @before_receive_args = args
+        end
+
+        on S2::Messages::Handshake do |message|
+          # noop
+        end
+      end
+
+      message = build(:s2_handshake)
+      connection = connection_class.new(ws, logger: Logger.new(nil))
+      connection.open("123")
+
+      expect { connection.receive_message(message.to_json) }
+        .to change { connection.before_receive_args }.from(nil).to(["123", message.to_json])
+    end
+
+    it "calls a callback after sending a message" do
+      connection_class = Class.new(described_class) do
+        attr_reader :after_send_args
+
+        after_send do |*args|
+          @after_send_args = args
+        end
+      end
+
+      allow(SecureRandom).to receive(:uuid).and_return("123")
+      message = build(:s2_handshake, message_id: "123")
+      payload = message.as_json.symbolize_keys
+
+      connection = connection_class.new(ws, logger: Logger.new(nil))
+      connection.open("123")
+
+      expect { connection.send_message(S2::Messages::Handshake, payload) }
+        .to change { connection.after_send_args }.from(nil).to(["123", message.to_json])
+    end
+  end
 end