From 3c227f356db744f6b07d3ebf7065fc593b7a0340 Mon Sep 17 00:00:00 2001
From: Martijn Versluis <martijnversluis@users.noreply.github.com>
Date: Thu, 12 Sep 2024 16:10:28 +0200
Subject: [PATCH] Implement Connection callbacks

Allows for defining callbacks in a connection implementation:

```ruby
on_open { |rm_id| puts rm_id }
before_receive { |rm_id, message_json| puts message_json }
after_send { |rm_id, message_json| puts message_json }
```
---
 lib/s2-ruby.rb                      |  1 +
 lib/s2/connection.rb                |  9 +++++
 lib/s2/message_handler_callbacks.rb | 39 +++++++++++++++++++
 spec/lib/s2/connection_spec.rb      | 58 +++++++++++++++++++++++++++++
 4 files changed, 107 insertions(+)
 create mode 100644 lib/s2/message_handler_callbacks.rb

diff --git a/lib/s2-ruby.rb b/lib/s2-ruby.rb
index abb9deb..1aab1c0 100644
--- a/lib/s2-ruby.rb
+++ b/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
diff --git a/lib/s2/connection.rb b/lib/s2/connection.rb
index bf97fa9..a88c0e2 100644
--- a/lib/s2/connection.rb
+++ b/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
diff --git a/lib/s2/message_handler_callbacks.rb b/lib/s2/message_handler_callbacks.rb
new file mode 100644
index 0000000..2d3705f
--- /dev/null
+++ b/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
diff --git a/spec/lib/s2/connection_spec.rb b/spec/lib/s2/connection_spec.rb
index 2ff0196..896e949 100644
--- a/spec/lib/s2/connection_spec.rb
+++ b/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