From cc9000a2acd64949d1cca841f709f6ff83da3ae8 Mon Sep 17 00:00:00 2001
From: Bob Forma <1178544+bforma@users.noreply.github.com>
Date: Thu, 12 Sep 2024 15:36:27 +0200
Subject: [PATCH] Expose state and update_state method

---
 lib/s2/connection.rb | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/lib/s2/connection.rb b/lib/s2/connection.rb
index 78bc3ab..bf97fa9 100644
--- a/lib/s2/connection.rb
+++ b/lib/s2/connection.rb
@@ -2,7 +2,7 @@ module S2
   class Connection
     include S2::MessageHandler
 
-    attr_reader :sent_messages
+    attr_reader :state, :sent_messages
 
     on S2::Messages::ReceptionStatus do |message|
       if message_sent?(message.subject_message_id)
@@ -16,8 +16,9 @@ class Connection
     def initialize(ws, logger: Rails.logger)
       @ws = ws
       @logger = logger
-      @state = :connected
       @sent_messages = {}
+
+      update_state :connected
     end
 
     def receive_message(message_json)
@@ -64,6 +65,10 @@ def reply(message, status:)
       end
     end
 
+    def update_state(new_state)
+      @state = new_state
+    end
+
     def message_sent?(message_id)
       @sent_messages.has_key?(message_id)
     end