From 2c56aef8d0d271a3ff4505ed53d8368043845b98 Mon Sep 17 00:00:00 2001
From: Roman Samoilov <2270393+rsamoilov@users.noreply.github.com>
Date: Tue, 30 Jul 2024 19:27:58 +0100
Subject: [PATCH] test

---
 spec/integration/integration_spec.rb |  3 +-
 spec/support/websocket_helper.rb     | 52 ++++++++++++++++++----------
 2 files changed, 36 insertions(+), 19 deletions(-)

diff --git a/spec/integration/integration_spec.rb b/spec/integration/integration_spec.rb
index 13ba1395..643a9a8c 100644
--- a/spec/integration/integration_spec.rb
+++ b/spec/integration/integration_spec.rb
@@ -336,10 +336,11 @@
         end
       end
 
-      sleep 1.5
+      sleep 2.5
 
       threads.each do |thread|
         client = thread.value
+        puts "[integration spec] #{client.messages}"
         expect(client.messages.last).to include("synced from time.com")
       end
     end
diff --git a/spec/support/websocket_helper.rb b/spec/support/websocket_helper.rb
index f080adcc..54c74d7a 100644
--- a/spec/support/websocket_helper.rb
+++ b/spec/support/websocket_helper.rb
@@ -11,11 +11,44 @@ def with_websocket_connection(url, headers: {})
 
   class WebSocketTestClient
     def initialize(url, headers: {})
+      @url = url
+      @headers = headers
+      connect
+    end
+
+    def connected?
+      @ws.handshake.valid? && @ws_data[:connected] && !@ws_data[:closed]
+    end
+
+    def send(data)
+      2.times do
+        @ws.send(data)
+        break
+      rescue Errno::EBADF
+        puts "[with_websocket_connection] reconnecting..."
+        @ws.connect(@url, headers: @headers)
+      end
+
+      sleep 0.1
+    end
+
+    def heartbeats
+      @ws_data[:heartbeats]
+    end
+
+    def messages
+      @ws_data[:messages]
+    end
+
+    private
+
+    def connect
       ws_data = { connected: false, closed: false, heartbeats: [], messages: [] }
 
-      @ws = WebSocket::Client::Simple.connect(url, headers:) do |ws|
+      @ws = WebSocket::Client::Simple.connect(@url, headers: @headers) do |ws|
         ws.on :open do
           ws_data[:connected] = true
+          ws_data[:closed] = false
         end
 
         ws.on :message do |msg|
@@ -32,22 +65,5 @@ def initialize(url, headers: {})
 
       sleep 0.1
     end
-
-    def connected?
-      @ws.handshake.valid? && @ws_data[:connected] && !@ws_data[:closed]
-    end
-
-    def send(data)
-      @ws.send(data)
-      sleep 0.1
-    end
-
-    def heartbeats
-      @ws_data[:heartbeats]
-    end
-
-    def messages
-      @ws_data[:messages]
-    end
   end
 end