Auto configure outbound websocket for sleepy shelly RPC devices (#124545)

This commit is contained in:
J. Nick Koston 2024-08-26 08:37:36 -10:00 committed by GitHub
parent b9aaba0432
commit 76182c246d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 259 additions and 22 deletions

View file

@ -254,3 +254,6 @@ VIRTUAL_NUMBER_MODE_MAP = {
"field": NumberMode.BOX,
"slider": NumberMode.SLIDER,
}
API_WS_URL = "/api/shelly/ws"

View file

@ -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,
)

View file

@ -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("/")))

View file

@ -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 = {

View file

@ -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: