From cd3059aa14e5972d829fd24b54bc6fb37777fcf7 Mon Sep 17 00:00:00 2001 From: IceBotYT <34712694+IceBotYT@users.noreply.github.com> Date: Fri, 6 Sep 2024 12:22:59 -0400 Subject: [PATCH] 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 --- homeassistant/components/nice_go/__init__.py | 8 +- .../components/nice_go/coordinator.py | 85 +++++++++-- homeassistant/components/nice_go/event.py | 6 +- .../components/nice_go/manifest.json | 3 +- tests/components/nice_go/test_diagnostics.py | 2 + tests/components/nice_go/test_event.py | 4 +- tests/components/nice_go/test_init.py | 141 ++++++++++++++++-- 7 files changed, 221 insertions(+), 28 deletions(-) diff --git a/homeassistant/components/nice_go/__init__.py b/homeassistant/components/nice_go/__init__.py index ab3dc06e3c1f09..b217112c192be0 100644 --- a/homeassistant/components/nice_go/__init__.py +++ b/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 diff --git a/homeassistant/components/nice_go/coordinator.py b/homeassistant/components/nice_go/coordinator.py index 323e0a08fe84ee..d6693db2d8a370 100644 --- a/homeassistant/components/nice_go/coordinator.py +++ b/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") diff --git a/homeassistant/components/nice_go/event.py b/homeassistant/components/nice_go/event.py index a19511b0b115c5..cd9198bcd26e29 100644 --- a/homeassistant/components/nice_go/event.py +++ b/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.""" diff --git a/homeassistant/components/nice_go/manifest.json b/homeassistant/components/nice_go/manifest.json index 884f2eb7b18ee4..315f23d949de1d 100644 --- a/homeassistant/components/nice_go/manifest.json +++ b/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"] } diff --git a/tests/components/nice_go/test_diagnostics.py b/tests/components/nice_go/test_diagnostics.py index f91f5748792032..5c8647f3d6e0d7 100644 --- a/tests/components/nice_go/test_diagnostics.py +++ b/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, diff --git a/tests/components/nice_go/test_event.py b/tests/components/nice_go/test_event.py index 0038b2882ad4cc..1c1b70532f4eb0 100644 --- a/tests/components/nice_go/test_event.py +++ b/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") diff --git a/tests/components/nice_go/test_init.py b/tests/components/nice_go/test_init.py index 5568a7ea62aa36..9c9bf28ca7a253 100644 --- a/tests/components/nice_go/test_init.py +++ b/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 homeassistant.components.nice_go.const import DOMAIN from homeassistant.config_entries import ConfigEntryState -from homeassistant.const import Platform -from homeassistant.core import HomeAssistant +from homeassistant.const import EVENT_HOMEASSISTANT_STOP, Platform +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 + )