Implement on_close callback

lib/s2/connection.rb

@@ -60,6 +60,10 @@ def close(message: nil)
       @ws.close
     end
 
+    def notify_closed(rm_id)
+      trigger_on_close(rm_id)
+    end
+
     protected
 
     def reply(message, status:)

lib/s2/message_handler_callbacks.rb

@@ -6,6 +6,7 @@ module MessageHandlerCallbacks
       class_attribute :on_open_proc, default: nil
       class_attribute :before_receive_proc, default: nil
       class_attribute :after_send_proc, default: nil
+      class_attribute :on_close_proc, default: nil
     end
 
     class_methods do
@@ -20,6 +21,10 @@ def before_receive(&block)
       def after_send(&block)
         self.after_send_proc = block
       end
+
+      def on_close(&block)
+        self.on_close_proc = block
+      end
     end
 
     protected
@@ -35,5 +40,9 @@ def trigger_before_receive(rm_id, message)
     def trigger_after_send(rm_id, message)
       instance_exec(rm_id, message, &self.class.after_send_proc) if self.class.after_send_proc
     end
+
+    def trigger_on_close(rm_id)
+      instance_exec(rm_id, &self.class.on_close_proc) if self.class.on_close_proc
+    end
   end
 end

spec/lib/s2/connection_spec.rb

@@ -159,5 +159,21 @@
       expect { connection.send_message(S2::Messages::Handshake, payload) }
         .to change { connection.after_send_args }.from(nil).to(["123", message.to_json])
     end
+
+    it "calls a callback on closing the connection" do
+      connection_class = Class.new(described_class) do
+        attr_reader :on_close_arg
+
+        on_close do |rm_id|
+          @on_close_arg = rm_id
+        end
+      end
+
+      connection = connection_class.new(ws, logger: Logger.new(nil))
+      connection.open("123")
+      connection.notify_closed("123")
+
+      expect(connection.on_close_arg).to eq("123")
+    end
   end
 end