fix: handle websocket disconnects better (#420)

* fix: handle websocket disconnects better * formatting * fix ws_start() call and fix some tests * catch assertion error if websocket isn't setup
firstof9 · Jan 29, 2025 · 4ca3485 · 4ca3485
1 parent 23bdfa2
commit 4ca3485
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 8 deletions.
diff --git a/custom_components/openevse/__init__.py b/custom_components/openevse/__init__.py
@@ -182,9 +182,6 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
     )
 
     await coordinator.async_refresh()
-    # Start the websocket listener
-    manager.ws_start()
-
     await hass.config_entries.async_forward_entry_setups(config_entry, PLATFORMS)
 
     services = OpenEVSEServices(hass, config_entry)
@@ -323,6 +320,26 @@ async def update_sensors(self) -> dict:
             )
             raise UpdateFailed(error) from error
 
+        try:
+            assert self._manager.websocket
+        except AssertionError as error:
+            _LOGGER.debug("Websocket not setup.")
+            raise UpdateFailed(error) from error
+
+        if self._manager.websocket.state != "connected":
+            _LOGGER.info("Connecting to websocket...")
+            try:
+                self._manager.ws_start()
+            except RuntimeError:
+                pass
+            except Exception as error:
+                _LOGGER.debug(
+                    "Error connecting to websocket [%s]: %s",
+                    type(error).__name__,
+                    error,
+                )
+                raise UpdateFailed(error) from error
+
         self.parse_sensors()
         await self.async_parse_sensors()
         _LOGGER.debug("Coordinator data: %s", self._data)

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -98,7 +98,7 @@ def test_charger(mock_aioclient):
     )
     mock_aioclient.get(
         TEST_URL_WS,
-        status=200,
+        status=101,
         body=load_fixture("status.json"),
         repeat=True,
     )
@@ -140,7 +140,7 @@ def test_charger_bad_serial(mock_aioclient):
     )
     mock_aioclient.get(
         TEST_URL_WS,
-        status=200,
+        status=101,
         body=load_fixture("status.json"),
         repeat=True,
     )
@@ -181,7 +181,7 @@ def test_charger_bad_post(mock_aioclient):
     )
     mock_aioclient.get(
         TEST_URL_WS,
-        status=200,
+        status=101,
         body=load_fixture("status.json"),
         repeat=True,
     )

diff --git a/tests/test_services.py b/tests/test_services.py
@@ -108,7 +108,7 @@ async def test_make_claim(
         data=CONFIG_DATA,
     )
     mock_aioclient.post(
-        f"{TEST_URL_CLAIMS}/4",
+        f"{TEST_URL_CLAIMS}/20",
         status=200,
         body='[{"msg":"done"}]',
         repeat=True,
@@ -150,7 +150,7 @@ async def test_release_claim(
         data=CONFIG_DATA,
     )
     mock_aioclient.delete(
-        f"{TEST_URL_CLAIMS}/4",
+        f"{TEST_URL_CLAIMS}/20",
         status=200,
         body='[{"msg":"done"}]',
         repeat=True,