From 61b23c4465595a8b37dc694a071bcbf30c438340 Mon Sep 17 00:00:00 2001
From: Daniel McKnight <daniel@neon.ai>
Date: Mon, 20 May 2024 15:46:08 -0700
Subject: [PATCH] Add WS disconnection handling Timeout if WS does not connect

---
 neon_nodes/websocket_client.py | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/neon_nodes/websocket_client.py b/neon_nodes/websocket_client.py
index 04a65c7..5c6dcdf 100644
--- a/neon_nodes/websocket_client.py
+++ b/neon_nodes/websocket_client.py
@@ -71,6 +71,11 @@ def __init__(self, bus=None, ready_hook=on_ready, error_hook=on_error,
         def ws_connect(*_, **__):
             self._connected.set()
 
+        def ws_disconnect(*_, **__):
+            self._connected.clear()
+            LOG.warning("Websocket disconnected")
+            self._wait_for_connection()
+
         def ws_error(_, exception):
             self.error_hook(exception)
             raise ConnectionError(f"Failed to connect: {exception}")
@@ -79,7 +84,8 @@ def ws_error(_, exception):
         self.websocket = WebSocketApp(f"{ws_address}/node/v1?token={auth_data['access_token']}",
                                       on_message=self._on_ws_data,
                                       on_open=ws_connect,
-                                      on_error=ws_error)
+                                      on_error=ws_error,
+                                      on_close=ws_disconnect)
         Thread(target=self.websocket.run_forever, daemon=True).start()
         self._device_data = self.config.get('neon_node', {})
         LOG.init(self.config.get("logging"))
@@ -112,10 +118,16 @@ def ws_error(_, exception):
 
         started_hook()
         self.run()
-        LOG.debug("Waiting for WS connection")
-        self._connected.wait()
+        self._wait_for_connection()
         ready_hook()
 
+    def _wait_for_connection(self):
+        LOG.debug("Waiting for WS connection")
+        if not self._connected.wait(30):
+            error = f"Timeout waiting for connection to {self.websocket.url}"
+            self.error_hook(error)
+            raise TimeoutError(error)
+
     @property
     def listening_sound(self) -> AudioSegment:
         """