Nice G.O. code quality improvements (#124319)

* Bring Nice G.O. up to platinum

* Switch to listen in coordinator

* Tests

* Remove parallel updates from coordinator

* Unsub from events on config entry unload

* Detect WS disconnection

* Tests

* Fix tests

* Set unsub to None after unsubbing

* Wait 5 seconds before setting update error to prevent excessive errors

* Tweaks

* More tweaks

* Tweaks part 2

* Potential test for hass stopping

* Improve reconnect handling and test on Homeassistant stop event

* Move event handler to entry init

* Patch const instead of asyncio.sleep

---------

Co-authored-by: jbouwh <[email protected]>

homeassistant/components/nice_go/__init__.py

@@ -5,7 +5,7 @@
 import logging
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import Platform
+from homeassistant.const import EVENT_HOMEASSISTANT_STOP, Platform
 from homeassistant.core import HomeAssistant
 
 from .coordinator import NiceGOUpdateCoordinator
@@ -25,8 +25,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: NiceGOConfigEntry) -> bo
     """Set up Nice G.O. from a config entry."""
 
     coordinator = NiceGOUpdateCoordinator(hass)
+    entry.async_on_unload(
+        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, coordinator.async_ha_stop)
+    )
 
     await coordinator.async_config_entry_first_refresh()
+
     entry.runtime_data = coordinator
 
     entry.async_create_background_task(
@@ -35,6 +39,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: NiceGOConfigEntry) -> bo
         "nice_go_websocket_task",
     )
 
+    entry.async_on_unload(coordinator.unsubscribe)
+
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True

homeassistant/components/nice_go/coordinator.py

@@ -3,11 +3,12 @@
 from __future__ import annotations
 
 import asyncio
+from collections.abc import Callable
 from dataclasses import dataclass
 from datetime import datetime
 import json
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from nice_go import (
     BARRIER_STATUS,
@@ -20,7 +21,7 @@
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_EMAIL, CONF_PASSWORD
-from homeassistant.core import HomeAssistant
+from homeassistant.core import Event, HomeAssistant, callback
 from homeassistant.exceptions import ConfigEntryAuthFailed
 from homeassistant.helpers import issue_registry as ir
 from homeassistant.helpers.aiohttp_client import async_get_clientsession
@@ -35,6 +36,9 @@
 
 _LOGGER = logging.getLogger(__name__)
 
+RECONNECT_ATTEMPTS = 3
+RECONNECT_DELAY = 5
+
 
 @dataclass
 class NiceGODevice:
@@ -70,7 +74,16 @@ def __init__(self, hass: HomeAssistant) -> None:
         self.email = self.config_entry.data[CONF_EMAIL]
         self.password = self.config_entry.data[CONF_PASSWORD]
         self.api = NiceGOApi()
-        self.ws_connected = False
+        self._unsub_connected: Callable[[], None] | None = None
+        self._unsub_data: Callable[[], None] | None = None
+        self._unsub_connection_lost: Callable[[], None] | None = None
+        self.connected = False
+        self._hass_stopping: bool = hass.is_stopping
+
+    @callback
+    def async_ha_stop(self, event: Event) -> None:
+        """Stop reconnecting if hass is stopping."""
+        self._hass_stopping = True
 
     async def _parse_barrier(self, barrier_state: BarrierState) -> NiceGODevice | None:
         """Parse barrier data."""
@@ -178,16 +191,30 @@ async def _update_refresh_token(self) -> None:
 
     async def client_listen(self) -> None:
         """Listen to the websocket for updates."""
-        self.api.event(self.on_connected)
-        self.api.event(self.on_data)
-        try:
-            await self.api.connect(reconnect=True)
-        except ApiError:
-            _LOGGER.exception("API error")
+        self._unsub_connected = self.api.listen("on_connected", self.on_connected)
+        self._unsub_data = self.api.listen("on_data", self.on_data)
+        self._unsub_connection_lost = self.api.listen(
+            "on_connection_lost", self.on_connection_lost
+        )
+
+        for _ in range(RECONNECT_ATTEMPTS):
+            if self._hass_stopping:
+                return
+
+            try:
+                await self.api.connect(reconnect=True)
+            except ApiError:
+                _LOGGER.exception("API error")
+            else:
+                return
+
+            await asyncio.sleep(RECONNECT_DELAY)
 
-        if not self.hass.is_stopping:
-            await asyncio.sleep(5)
-            await self.client_listen()
+        self.async_set_update_error(
+            TimeoutError(
+                "Failed to connect to the websocket, reconnect attempts exhausted"
+            )
+        )
 
     async def on_data(self, data: dict[str, Any]) -> None:
         """Handle incoming data from the websocket."""
@@ -220,4 +247,38 @@ async def on_data(self, data: dict[str, Any]) -> None:
     async def on_connected(self) -> None:
         """Handle the websocket connection."""
         _LOGGER.debug("Connected to the websocket")
+        self.connected = True
+
         await self.api.subscribe(self.organization_id)
+
+        if not self.last_update_success:
+            self.async_set_updated_data(self.data)
+
+    async def on_connection_lost(self, data: dict[str, Exception]) -> None:
+        """Handle the websocket connection loss. Don't need to do much since the library will automatically reconnect."""
+        _LOGGER.debug("Connection lost to the websocket")
+        self.connected = False
+
+        # Give some time for reconnection
+        await asyncio.sleep(RECONNECT_DELAY)
+        if self.connected:
+            _LOGGER.debug("Reconnected, not setting error")
+            return
+
+        # There's likely a problem with the connection, and not the server being flaky
+        self.async_set_update_error(data["exception"])
+
+    def unsubscribe(self) -> None:
+        """Unsubscribe from the websocket."""
+        if TYPE_CHECKING:
+            assert self._unsub_connected is not None
+            assert self._unsub_data is not None
+            assert self._unsub_connection_lost is not None
+
+        self._unsub_connection_lost()
+        self._unsub_connected()
+        self._unsub_data()
+        self._unsub_connected = None
+        self._unsub_data = None
+        self._unsub_connection_lost = None
+        _LOGGER.debug("Unsubscribed from the websocket")

homeassistant/components/nice_go/event.py

@@ -40,7 +40,11 @@ class NiceGOEventEntity(NiceGOEntity, EventEntity):
     async def async_added_to_hass(self) -> None:
         """Listen for events."""
         await super().async_added_to_hass()
-        self.coordinator.api.event(self.on_barrier_obstructed)
+        self.async_on_remove(
+            self.coordinator.api.listen(
+                "on_barrier_obstructed", self.on_barrier_obstructed
+            )
+        )
 
     async def on_barrier_obstructed(self, data: dict[str, Any]) -> None:
         """Handle barrier obstructed event."""

homeassistant/components/nice_go/manifest.json

@@ -4,7 +4,8 @@
   "codeowners": ["@IceBotYT"],
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/nice_go",
+  "integration_type": "hub",
   "iot_class": "cloud_push",
-  "loggers": ["nice-go"],
+  "loggers": ["nice_go"],
   "requirements": ["nice-go==0.3.8"]
 }

tests/components/nice_go/test_diagnostics.py

@@ -2,6 +2,7 @@
 
 from unittest.mock import AsyncMock
 
+import pytest
 from syrupy import SnapshotAssertion
 from syrupy.filters import props
 
@@ -14,6 +15,7 @@
 from tests.typing import ClientSessionGenerator
 
 
+@pytest.mark.freeze_time("2024-08-27")
 async def test_entry_diagnostics(
     hass: HomeAssistant,
     hass_client: ClientSessionGenerator,

tests/components/nice_go/test_event.py

@@ -19,10 +19,10 @@ async def test_barrier_obstructed(
     mock_config_entry: MockConfigEntry,
 ) -> None:
     """Test barrier obstructed."""
-    mock_nice_go.event = MagicMock()
+    mock_nice_go.listen = MagicMock()
     await setup_integration(hass, mock_config_entry, [Platform.EVENT])
 
-    await mock_nice_go.event.call_args_list[2][0][0]({"deviceId": "1"})
+    await mock_nice_go.listen.call_args_list[3][0][1]({"deviceId": "1"})
     await hass.async_block_till_done()
 
     event_state = hass.states.get("event.test_garage_1_barrier_obstructed")