From 76182c246d8f2c302d1b85f4443aa3d00b9bb9cc Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Mon, 26 Aug 2024 08:37:36 -1000 Subject: [PATCH] Auto configure outbound websocket for sleepy shelly RPC devices (#124545) --- homeassistant/components/shelly/const.py | 3 + .../components/shelly/coordinator.py | 82 ++++++-- homeassistant/components/shelly/utils.py | 17 +- tests/components/shelly/conftest.py | 1 + tests/components/shelly/test_config_flow.py | 178 +++++++++++++++++- 5 files changed, 259 insertions(+), 22 deletions(-) diff --git a/homeassistant/components/shelly/const.py b/homeassistant/components/shelly/const.py index 1759f4bdd18..fe4108a1f52 100644 --- a/homeassistant/components/shelly/const.py +++ b/homeassistant/components/shelly/const.py @@ -254,3 +254,6 @@ VIRTUAL_NUMBER_MODE_MAP = { "field": NumberMode.BOX, "slider": NumberMode.SLIDER, } + + +API_WS_URL = "/api/shelly/ws" diff --git a/homeassistant/components/shelly/coordinator.py b/homeassistant/components/shelly/coordinator.py index 6286e515727..03dcdedbb6f 100644 --- a/homeassistant/components/shelly/coordinator.py +++ b/homeassistant/components/shelly/coordinator.py @@ -64,6 +64,7 @@ from .utils import ( get_host, get_http_port, get_rpc_device_wakeup_period, + get_rpc_ws_url, update_device_fw_info, ) @@ -101,6 +102,9 @@ class ShellyCoordinatorBase[_DeviceT: BlockDevice | RpcDevice]( self._pending_platforms: list[Platform] | None = None device_name = device.name if device.initialized else entry.title interval_td = timedelta(seconds=update_interval) + # The device has come online at least once. In the case of a sleeping RPC + # device, this means that the device has connected to the WS server at least once. + self._came_online_once = False super().__init__(hass, LOGGER, name=device_name, update_interval=interval_td) self._debounced_reload: Debouncer[Coroutine[Any, Any, None]] = Debouncer( @@ -184,7 +188,7 @@ class ShellyCoordinatorBase[_DeviceT: BlockDevice | RpcDevice]( if not self._pending_platforms: return True - LOGGER.debug("Device %s is online, resuming setup", self.entry.title) + LOGGER.debug("Device %s is online, resuming setup", self.name) platforms = self._pending_platforms self._pending_platforms = None @@ -372,6 +376,7 @@ class ShellyBlockCoordinator(ShellyCoordinatorBase[BlockDevice]): """Handle device update.""" LOGGER.debug("Shelly %s handle update, type: %s", self.name, update_type) if update_type is BlockUpdateType.ONLINE: + self._came_online_once = True self.entry.async_create_background_task( self.hass, self._async_device_connect_task(), @@ -472,9 +477,24 @@ class ShellyRpcCoordinator(ShellyCoordinatorBase[RpcDevice]): self._event_listeners: list[Callable[[dict[str, Any]], None]] = [] self._ota_event_listeners: list[Callable[[dict[str, Any]], None]] = [] self._input_event_listeners: list[Callable[[dict[str, Any]], None]] = [] - + self._connect_task: asyncio.Task | None = None entry.async_on_unload(entry.add_update_listener(self._async_update_listener)) + async def async_device_online(self) -> None: + """Handle device going online.""" + if not self.sleep_period: + await self.async_request_refresh() + elif not self._came_online_once or not self.device.initialized: + LOGGER.debug( + "Sleepy device %s is online, trying to poll and configure", self.name + ) + # Zeroconf told us the device is online, try to poll + # the device and if possible, set up the outbound + # websocket so the device will send us updates + # instead of relying on polling it fast enough before + # it goes to sleep again + self._async_handle_rpc_device_online() + def update_sleep_period(self) -> bool: """Check device sleep period & update if changed.""" if ( @@ -598,15 +618,15 @@ class ShellyRpcCoordinator(ShellyCoordinatorBase[RpcDevice]): async def _async_disconnected(self, reconnect: bool) -> None: """Handle device disconnected.""" - # Sleeping devices send data and disconnect - # There are no disconnect events for sleeping devices - if self.sleep_period: - return - async with self._connection_lock: if not self.connected: # Already disconnected return self.connected = False + # Sleeping devices send data and disconnect + # There are no disconnect events for sleeping devices + # but we do need to make sure self.connected is False + if self.sleep_period: + return self._async_run_disconnected_events() # Try to reconnect right away if triggered by disconnect event if reconnect: @@ -645,6 +665,21 @@ class ShellyRpcCoordinator(ShellyCoordinatorBase[RpcDevice]): """ if not self.sleep_period: await self._async_connect_ble_scanner() + else: + await self._async_setup_outbound_websocket() + + async def _async_setup_outbound_websocket(self) -> None: + """Set up outbound websocket if it is not enabled.""" + config = self.device.config + if ( + (ws_config := config.get("ws")) + and (not ws_config["server"] or not ws_config["enable"]) + and (ws_url := get_rpc_ws_url(self.hass)) + ): + LOGGER.debug( + "Setting up outbound websocket for device %s - %s", self.name, ws_url + ) + await self.device.update_outbound_websocket(ws_url) async def _async_connect_ble_scanner(self) -> None: """Connect BLE scanner.""" @@ -662,6 +697,21 @@ class ShellyRpcCoordinator(ShellyCoordinatorBase[RpcDevice]): await async_connect_scanner(self.hass, self, ble_scanner_mode) ) + @callback + def _async_handle_rpc_device_online(self) -> None: + """Handle device going online.""" + if self.device.connected or ( + self._connect_task and not self._connect_task.done() + ): + LOGGER.debug("Device %s already connected/connecting", self.name) + return + self._connect_task = self.entry.async_create_background_task( + self.hass, + self._async_device_connect_task(), + "rpc device online", + eager_start=True, + ) + @callback def _async_handle_update( self, device_: RpcDevice, update_type: RpcUpdateType @@ -669,15 +719,8 @@ class ShellyRpcCoordinator(ShellyCoordinatorBase[RpcDevice]): """Handle device update.""" LOGGER.debug("Shelly %s handle update, type: %s", self.name, update_type) if update_type is RpcUpdateType.ONLINE: - if self.device.connected: - LOGGER.debug("Device %s already connected", self.name) - return - self.entry.async_create_background_task( - self.hass, - self._async_device_connect_task(), - "rpc device online", - eager_start=True, - ) + self._came_online_once = True + self._async_handle_rpc_device_online() elif update_type is RpcUpdateType.INITIALIZED: self.entry.async_create_background_task( self.hass, self._async_connected(), "rpc device init", eager_start=True @@ -798,14 +841,13 @@ def get_rpc_coordinator_by_device_id( async def async_reconnect_soon(hass: HomeAssistant, entry: ShellyConfigEntry) -> None: """Try to reconnect soon.""" if ( - not entry.data.get(CONF_SLEEP_PERIOD) - and not hass.is_stopping - and entry.state == ConfigEntryState.LOADED + not hass.is_stopping + and entry.state is ConfigEntryState.LOADED and (coordinator := entry.runtime_data.rpc) ): entry.async_create_background_task( hass, - coordinator.async_request_refresh(), + coordinator.async_device_online(), "reconnect soon", eager_start=True, ) diff --git a/homeassistant/components/shelly/utils.py b/homeassistant/components/shelly/utils.py index 339f6781171..d0a8a1230c5 100644 --- a/homeassistant/components/shelly/utils.py +++ b/homeassistant/components/shelly/utils.py @@ -23,6 +23,7 @@ from aioshelly.const import ( RPC_GENERATIONS, ) from aioshelly.rpc_device import RpcDevice, WsServer +from yarl import URL from homeassistant.components import network from homeassistant.components.http import HomeAssistantView @@ -36,9 +37,11 @@ from homeassistant.helpers import ( singleton, ) from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC +from homeassistant.helpers.network import NoURLAvailableError, get_url from homeassistant.util.dt import utcnow from .const import ( + API_WS_URL, BASIC_INPUTS_EVENTS_TYPES, CONF_COAP_PORT, CONF_GEN, @@ -254,7 +257,7 @@ class ShellyReceiver(HomeAssistantView): """Handle pushes from Shelly Gen2 devices.""" requires_auth = False - url = "/api/shelly/ws" + url = API_WS_URL name = "api:shelly:ws" def __init__(self, ws_server: WsServer) -> None: @@ -571,3 +574,15 @@ def async_remove_orphaned_virtual_entities( if orphaned_entities: async_remove_shelly_rpc_entities(hass, platform, mac, orphaned_entities) + + +def get_rpc_ws_url(hass: HomeAssistant) -> str | None: + """Return the RPC websocket URL.""" + try: + raw_url = get_url(hass, prefer_external=False, allow_cloud=False) + except NoURLAvailableError: + LOGGER.debug("URL not available, skipping outbound websocket setup") + return None + url = URL(raw_url) + ws_url = url.with_scheme("wss" if url.scheme == "https" else "ws") + return str(ws_url.joinpath(API_WS_URL.removeprefix("/"))) diff --git a/tests/components/shelly/conftest.py b/tests/components/shelly/conftest.py index a2629d21362..34e4ce1379e 100644 --- a/tests/components/shelly/conftest.py +++ b/tests/components/shelly/conftest.py @@ -186,6 +186,7 @@ MOCK_CONFIG = { "device": {"name": "Test name"}, }, "wifi": {"sta": {"enable": True}, "sta1": {"enable": False}}, + "ws": {"enable": False, "server": None}, } MOCK_SHELLY_COAP = { diff --git a/tests/components/shelly/test_config_flow.py b/tests/components/shelly/test_config_flow.py index 0c574a33e0c..c0c089f469a 100644 --- a/tests/components/shelly/test_config_flow.py +++ b/tests/components/shelly/test_config_flow.py @@ -4,7 +4,7 @@ from dataclasses import replace from datetime import timedelta from ipaddress import ip_address from typing import Any -from unittest.mock import AsyncMock, Mock, patch +from unittest.mock import AsyncMock, Mock, call, patch from aioshelly.const import DEFAULT_HTTP_PORT, MODEL_1, MODEL_PLUS_2PM from aioshelly.exceptions import ( @@ -1153,6 +1153,182 @@ async def test_zeroconf_sleeping_device_not_triggers_refresh( assert "device did not update" not in caplog.text +async def test_zeroconf_sleeping_device_attempts_configure( + hass: HomeAssistant, + mock_rpc_device: Mock, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test zeroconf discovery configures a sleeping device outbound websocket.""" + monkeypatch.setattr(mock_rpc_device, "connected", False) + monkeypatch.setattr(mock_rpc_device, "initialized", False) + monkeypatch.setitem(mock_rpc_device.status["sys"], "wakeup_period", 1000) + entry = MockConfigEntry( + domain="shelly", + unique_id="AABBCCDDEEFF", + data={"host": "1.1.1.1", "gen": 2, "sleep_period": 1000, "model": MODEL_1}, + ) + entry.add_to_hass(hass) + await hass.config_entries.async_setup(entry.entry_id) + await hass.async_block_till_done() + mock_rpc_device.mock_disconnected() + await hass.async_block_till_done() + + mock_rpc_device.mock_online() + await hass.async_block_till_done(wait_background_tasks=True) + + assert "online, resuming setup" in caplog.text + assert len(mock_rpc_device.initialize.mock_calls) == 1 + + with patch( + "homeassistant.components.shelly.config_flow.get_info", + return_value={"mac": "AABBCCDDEEFF", "type": MODEL_1, "auth": False}, + ): + result = await hass.config_entries.flow.async_init( + DOMAIN, + data=DISCOVERY_INFO, + context={"source": config_entries.SOURCE_ZEROCONF}, + ) + assert result["type"] is FlowResultType.ABORT + assert result["reason"] == "already_configured" + + assert mock_rpc_device.update_outbound_websocket.mock_calls == [] + + monkeypatch.setattr(mock_rpc_device, "connected", True) + monkeypatch.setattr(mock_rpc_device, "initialized", True) + mock_rpc_device.mock_initialized() + async_fire_time_changed( + hass, dt_util.utcnow() + timedelta(seconds=ENTRY_RELOAD_COOLDOWN) + ) + await hass.async_block_till_done() + assert "device did not update" not in caplog.text + + monkeypatch.setattr(mock_rpc_device, "connected", False) + mock_rpc_device.mock_disconnected() + assert mock_rpc_device.update_outbound_websocket.mock_calls == [ + call("ws://10.10.10.10:8123/api/shelly/ws") + ] + + +async def test_zeroconf_sleeping_device_attempts_configure_ws_disabled( + hass: HomeAssistant, + mock_rpc_device: Mock, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test zeroconf discovery configures a sleeping device outbound websocket when its disabled.""" + monkeypatch.setattr(mock_rpc_device, "connected", False) + monkeypatch.setattr(mock_rpc_device, "initialized", False) + monkeypatch.setitem(mock_rpc_device.status["sys"], "wakeup_period", 1000) + monkeypatch.setitem( + mock_rpc_device.config, "ws", {"enable": False, "server": "ws://oldha"} + ) + entry = MockConfigEntry( + domain="shelly", + unique_id="AABBCCDDEEFF", + data={"host": "1.1.1.1", "gen": 2, "sleep_period": 1000, "model": MODEL_1}, + ) + entry.add_to_hass(hass) + await hass.config_entries.async_setup(entry.entry_id) + await hass.async_block_till_done() + mock_rpc_device.mock_disconnected() + await hass.async_block_till_done() + + mock_rpc_device.mock_online() + await hass.async_block_till_done(wait_background_tasks=True) + + assert "online, resuming setup" in caplog.text + assert len(mock_rpc_device.initialize.mock_calls) == 1 + + with patch( + "homeassistant.components.shelly.config_flow.get_info", + return_value={"mac": "AABBCCDDEEFF", "type": MODEL_1, "auth": False}, + ): + result = await hass.config_entries.flow.async_init( + DOMAIN, + data=DISCOVERY_INFO, + context={"source": config_entries.SOURCE_ZEROCONF}, + ) + assert result["type"] is FlowResultType.ABORT + assert result["reason"] == "already_configured" + + assert mock_rpc_device.update_outbound_websocket.mock_calls == [] + + monkeypatch.setattr(mock_rpc_device, "connected", True) + monkeypatch.setattr(mock_rpc_device, "initialized", True) + mock_rpc_device.mock_initialized() + async_fire_time_changed( + hass, dt_util.utcnow() + timedelta(seconds=ENTRY_RELOAD_COOLDOWN) + ) + await hass.async_block_till_done() + assert "device did not update" not in caplog.text + + monkeypatch.setattr(mock_rpc_device, "connected", False) + mock_rpc_device.mock_disconnected() + assert mock_rpc_device.update_outbound_websocket.mock_calls == [ + call("ws://10.10.10.10:8123/api/shelly/ws") + ] + + +async def test_zeroconf_sleeping_device_attempts_configure_no_url_available( + hass: HomeAssistant, + mock_rpc_device: Mock, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + """Test zeroconf discovery for sleeping device with no hass url.""" + hass.config.internal_url = None + hass.config.external_url = None + hass.config.api = None + monkeypatch.setattr(mock_rpc_device, "connected", False) + monkeypatch.setattr(mock_rpc_device, "initialized", False) + monkeypatch.setitem(mock_rpc_device.status["sys"], "wakeup_period", 1000) + entry = MockConfigEntry( + domain="shelly", + unique_id="AABBCCDDEEFF", + data={"host": "1.1.1.1", "gen": 2, "sleep_period": 1000, "model": MODEL_1}, + ) + entry.add_to_hass(hass) + await hass.config_entries.async_setup(entry.entry_id) + await hass.async_block_till_done() + mock_rpc_device.mock_disconnected() + await hass.async_block_till_done() + + mock_rpc_device.mock_online() + await hass.async_block_till_done(wait_background_tasks=True) + + assert "online, resuming setup" in caplog.text + assert len(mock_rpc_device.initialize.mock_calls) == 1 + + with patch( + "homeassistant.components.shelly.config_flow.get_info", + return_value={"mac": "AABBCCDDEEFF", "type": MODEL_1, "auth": False}, + ): + result = await hass.config_entries.flow.async_init( + DOMAIN, + data=DISCOVERY_INFO, + context={"source": config_entries.SOURCE_ZEROCONF}, + ) + assert result["type"] is FlowResultType.ABORT + assert result["reason"] == "already_configured" + + assert mock_rpc_device.update_outbound_websocket.mock_calls == [] + + monkeypatch.setattr(mock_rpc_device, "connected", True) + monkeypatch.setattr(mock_rpc_device, "initialized", True) + mock_rpc_device.mock_initialized() + async_fire_time_changed( + hass, dt_util.utcnow() + timedelta(seconds=ENTRY_RELOAD_COOLDOWN) + ) + await hass.async_block_till_done() + assert "device did not update" not in caplog.text + + monkeypatch.setattr(mock_rpc_device, "connected", False) + mock_rpc_device.mock_disconnected() + # No url available so no attempt to configure the device + assert mock_rpc_device.update_outbound_websocket.mock_calls == [] + + async def test_sleeping_device_gen2_with_new_firmware( hass: HomeAssistant, mock_rpc_device: Mock, monkeypatch: pytest.MonkeyPatch ) -> None: