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 <jan@jbsoft.nl>

homeassistant/components/nice_go/__init__.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 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 nice_go import (
 
 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 @@ from .const import (
 
 _LOGGER = logging.getLogger(__name__)
 
+RECONNECT_ATTEMPTS = 3
+RECONNECT_DELAY = 5
+
 
 @dataclass
 class NiceGODevice:
@@ -70,7 +74,16 @@ class NiceGOUpdateCoordinator(DataUpdateCoordinator[dict[str, NiceGODevice]]):
         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 @@ class NiceGOUpdateCoordinator(DataUpdateCoordinator[dict[str, NiceGODevice]]):
 
     async def client_listen(self) -> None:
         """Listen to the websocket for updates."""
-        self.api.event(self.on_connected)
+        self._unsub_connected = self.api.listen("on_connected", self.on_connected)
-        self.api.event(self.on_data)
+        self._unsub_data = self.api.listen("on_data", self.on_data)
-        try:
+        self._unsub_connection_lost = self.api.listen(
-            await self.api.connect(reconnect=True)
+            "on_connection_lost", self.on_connection_lost
-        except ApiError:
+        )
-            _LOGGER.exception("API error")
 
-        if not self.hass.is_stopping:
+        for _ in range(RECONNECT_ATTEMPTS):
-            await asyncio.sleep(5)
+            if self._hass_stopping:
-            await self.client_listen()
+                return
+
+            try:
+                await self.api.connect(reconnect=True)
+            except ApiError:
+                _LOGGER.exception("API error")
+            else:
+                return
+
+            await asyncio.sleep(RECONNECT_DELAY)
+
+        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 @@ class NiceGOUpdateCoordinator(DataUpdateCoordinator[dict[str, NiceGODevice]]):
     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.components.diagnostics import get_diagnostics_for_config_entry
 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")

tests/components/nice_go/test_init.py

@@ -1,7 +1,8 @@
 """Test Nice G.O. init."""
 
+import asyncio
 from datetime import timedelta
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock, patch
 
 from freezegun.api import FrozenDateTimeFactory
 from nice_go import ApiError, AuthFailedError, Barrier, BarrierState
@@ -10,8 +11,8 @@ from syrupy.assertion import SnapshotAssertion
 
 from homeassistant.components.nice_go.const import DOMAIN
 from homeassistant.config_entries import ConfigEntryState
-from homeassistant.const import Platform
+from homeassistant.const import EVENT_HOMEASSISTANT_STOP, Platform
-from homeassistant.core import HomeAssistant
+from homeassistant.core import Event, HomeAssistant, callback
 from homeassistant.helpers import issue_registry as ir
 
 from . import setup_integration
@@ -209,11 +210,11 @@ async def test_on_data_none_parsed(
 ) -> None:
     """Test on data with None parsed."""
 
-    mock_nice_go.event = MagicMock()
+    mock_nice_go.listen = MagicMock()
 
     await setup_integration(hass, mock_config_entry, [Platform.COVER])
 
-    await mock_nice_go.event.call_args[0][0](
+    await mock_nice_go.listen.call_args_list[1][0][1](
         {
             "data": {
                 "devicesStatesUpdateFeed": {
@@ -243,18 +244,74 @@ async def test_on_connected(
 ) -> None:
     """Test on connected."""
 
-    mock_nice_go.event = MagicMock()
+    mock_nice_go.listen = MagicMock()
 
     await setup_integration(hass, mock_config_entry, [Platform.COVER])
 
-    assert mock_nice_go.event.call_count == 2
+    assert mock_nice_go.listen.call_count == 3
 
     mock_nice_go.subscribe = AsyncMock()
-    await mock_nice_go.event.call_args_list[0][0][0]()
+    await mock_nice_go.listen.call_args_list[0][0][1]()
 
     assert mock_nice_go.subscribe.call_count == 1
 
 
+async def test_on_connection_lost(
+    hass: HomeAssistant,
+    mock_nice_go: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+    freezer: FrozenDateTimeFactory,
+) -> None:
+    """Test on connection lost."""
+
+    mock_nice_go.listen = MagicMock()
+
+    await setup_integration(hass, mock_config_entry, [Platform.COVER])
+
+    assert mock_nice_go.listen.call_count == 3
+
+    with patch("homeassistant.components.nice_go.coordinator.RECONNECT_DELAY", 0):
+        await mock_nice_go.listen.call_args_list[2][0][1](
+            {"exception": ValueError("test")}
+        )
+
+    assert hass.states.get("cover.test_garage_1").state == "unavailable"
+
+    # Now fire connected
+
+    mock_nice_go.subscribe = AsyncMock()
+
+    await mock_nice_go.listen.call_args_list[0][0][1]()
+
+    assert mock_nice_go.subscribe.call_count == 1
+
+    assert hass.states.get("cover.test_garage_1").state == "closed"
+
+
+async def test_on_connection_lost_reconnect(
+    hass: HomeAssistant,
+    mock_nice_go: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+    freezer: FrozenDateTimeFactory,
+) -> None:
+    """Test on connection lost with reconnect."""
+
+    mock_nice_go.listen = MagicMock()
+
+    await setup_integration(hass, mock_config_entry, [Platform.COVER])
+
+    assert mock_nice_go.listen.call_count == 3
+
+    assert hass.states.get("cover.test_garage_1").state == "closed"
+
+    with patch("homeassistant.components.nice_go.coordinator.RECONNECT_DELAY", 0):
+        await mock_nice_go.listen.call_args_list[2][0][1](
+            {"exception": ValueError("test")}
+        )
+
+    assert hass.states.get("cover.test_garage_1").state == "unavailable"
+
+
 async def test_no_connection_state(
     hass: HomeAssistant,
     mock_nice_go: AsyncMock,
@@ -262,13 +319,13 @@ async def test_no_connection_state(
 ) -> None:
     """Test parsing barrier with no connection state."""
 
-    mock_nice_go.event = MagicMock()
+    mock_nice_go.listen = MagicMock()
 
     await setup_integration(hass, mock_config_entry, [Platform.COVER])
 
-    assert mock_nice_go.event.call_count == 2
+    assert mock_nice_go.listen.call_count == 3
 
-    await mock_nice_go.event.call_args[0][0](
+    await mock_nice_go.listen.call_args_list[1][0][1](
         {
             "data": {
                 "devicesStatesUpdateFeed": {
@@ -286,3 +343,65 @@ async def test_no_connection_state(
     )
 
     assert hass.states.get("cover.test_garage_1").state == "unavailable"
+
+
+async def test_connection_attempts_exhausted(
+    hass: HomeAssistant,
+    mock_nice_go: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+    freezer: FrozenDateTimeFactory,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test connection attempts exhausted."""
+
+    mock_nice_go.connect.side_effect = ApiError
+
+    with (
+        patch("homeassistant.components.nice_go.coordinator.RECONNECT_ATTEMPTS", 1),
+        patch("homeassistant.components.nice_go.coordinator.RECONNECT_DELAY", 0),
+    ):
+        await setup_integration(hass, mock_config_entry, [Platform.COVER])
+
+    assert "API error" in caplog.text
+    assert "Error requesting Nice G.O. data" in caplog.text
+
+
+async def test_reconnect_hass_stopping(
+    hass: HomeAssistant,
+    mock_nice_go: AsyncMock,
+    mock_config_entry: MockConfigEntry,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test reconnect with hass stopping."""
+
+    mock_nice_go.listen = MagicMock()
+    mock_nice_go.connect.side_effect = ApiError
+
+    wait_for_hass = asyncio.Event()
+
+    @callback
+    def _async_ha_stop(event: Event) -> None:
+        """Stop reconnecting if hass is stopping."""
+        wait_for_hass.set()
+
+    hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, _async_ha_stop)
+
+    with (
+        patch("homeassistant.components.nice_go.coordinator.RECONNECT_DELAY", 0.1),
+        patch("homeassistant.components.nice_go.coordinator.RECONNECT_ATTEMPTS", 20),
+    ):
+        await setup_integration(hass, mock_config_entry, [Platform.COVER])
+        await hass.async_block_till_done()
+        hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP)
+        await wait_for_hass.wait()
+        await hass.async_block_till_done(wait_background_tasks=True)
+
+        assert mock_nice_go.connect.call_count < 10
+
+        assert len(hass._background_tasks) == 0
+
+        assert "API error" in caplog.text
+        assert (
+            "Failed to connect to the websocket, reconnect attempts exhausted"
+            not in caplog.text
+        )