diff --git a/supervisor/const.py b/supervisor/const.py index 3b36be2a3..8f77a9c98 100644 --- a/supervisor/const.py +++ b/supervisor/const.py @@ -451,6 +451,7 @@ class BusEvent(str, Enum): HARDWARE_NEW_DEVICE = "hardware_new_device" HARDWARE_REMOVE_DEVICE = "hardware_remove_device" DOCKER_CONTAINER_STATE_CHANGE = "docker_container_state_change" + SUPERVISOR_STATE_CHANGE = "supervisor_state_change" class CpuArch(str, Enum): @@ -461,3 +462,10 @@ class CpuArch(str, Enum): AARCH64 = "aarch64" I386 = "i386" AMD64 = "amd64" + + +STARTING_STATES = [ + CoreState.INITIALIZE, + CoreState.STARTUP, + CoreState.SETUP, +] diff --git a/supervisor/core.py b/supervisor/core.py index f8e62321c..cbdccaeb3 100644 --- a/supervisor/core.py +++ b/supervisor/core.py @@ -7,7 +7,14 @@ import logging import async_timeout -from .const import RUN_SUPERVISOR_STATE, AddonStartup, CoreState +from .const import ( + ATTR_STARTUP, + RUN_SUPERVISOR_STATE, + STARTING_STATES, + AddonStartup, + BusEvent, + CoreState, +) from .coresys import CoreSys, CoreSysAttributes from .exceptions import ( HassioError, @@ -63,9 +70,13 @@ class Core(CoreSysAttributes): ) finally: self._state = new_state - self.sys_homeassistant.websocket.supervisor_update_event( - "info", {"state": new_state} - ) + self.sys_bus.fire_event(BusEvent.SUPERVISOR_STATE_CHANGE, new_state) + + # These will be received by HA after startup has completed which won't make sense + if new_state not in STARTING_STATES: + self.sys_homeassistant.websocket.supervisor_update_event( + "info", {"state": new_state} + ) async def connect(self): """Connect Supervisor container.""" @@ -266,7 +277,9 @@ class Core(CoreSysAttributes): self.sys_create_task(self.sys_resolution.healthcheck()) self.state = CoreState.RUNNING - self.sys_homeassistant.websocket.supervisor_update_event("supervisor", {}) + self.sys_homeassistant.websocket.supervisor_update_event( + "supervisor", {ATTR_STARTUP: "complete"} + ) _LOGGER.info("Supervisor is up and running") async def stop(self): diff --git a/supervisor/homeassistant/module.py b/supervisor/homeassistant/module.py index 0a9bf2077..f85321f0d 100644 --- a/supervisor/homeassistant/module.py +++ b/supervisor/homeassistant/module.py @@ -260,6 +260,7 @@ class HomeAssistant(FileConfiguration, CoreSysAttributes): """Prepare Home Assistant object.""" await asyncio.wait( [ + self.sys_create_task(self.websocket.load()), self.sys_create_task(self.secrets.load()), self.sys_create_task(self.core.load()), ] diff --git a/supervisor/homeassistant/websocket.py b/supervisor/homeassistant/websocket.py index 79ff3af02..b0d355ace 100644 --- a/supervisor/homeassistant/websocket.py +++ b/supervisor/homeassistant/websocket.py @@ -9,7 +9,16 @@ import aiohttp from aiohttp.http_websocket import WSMsgType from awesomeversion import AwesomeVersion -from ..const import ATTR_ACCESS_TOKEN, ATTR_DATA, ATTR_EVENT, ATTR_TYPE, ATTR_UPDATE_KEY +from ..const import ( + ATTR_ACCESS_TOKEN, + ATTR_DATA, + ATTR_EVENT, + ATTR_TYPE, + ATTR_UPDATE_KEY, + STARTING_STATES, + BusEvent, + CoreState, +) from ..coresys import CoreSys, CoreSysAttributes from ..exceptions import ( HomeAssistantAPIError, @@ -172,6 +181,15 @@ class HomeAssistantWebSocket(CoreSysAttributes): self.coresys: CoreSys = coresys self._client: WSClient | None = None self._lock: asyncio.Lock = asyncio.Lock() + self._queue: list[dict[str, Any]] = [] + + async def _process_queue(self, reference: CoreState) -> None: + """Process queue once supervisor is running.""" + if reference == CoreState.RUNNING: + for msg in self._queue: + await self.async_send_message(msg) + + self._queue.clear() async def _get_ws_client(self) -> WSClient: """Return a websocket client.""" @@ -219,8 +237,21 @@ class HomeAssistantWebSocket(CoreSysAttributes): return False return True + async def load(self) -> None: + """Set up queue processor after startup completes.""" + self.sys_bus.register_event( + BusEvent.SUPERVISOR_STATE_CHANGE, self._process_queue + ) + async def async_send_message(self, message: dict[str, Any]) -> None: """Send a command with the WS client.""" + # Only commands allowed during startup as those tell Home Assistant to do something. + # Messages may cause clients to make follow-up API calls so those wait. + if self.sys_core.state in STARTING_STATES: + self._queue.append(message) + _LOGGER.debug("Queuing message until startup has completed: %s", message) + return + if not await self._can_send(message): return diff --git a/tests/homeassistant/test_module.py b/tests/homeassistant/test_module.py index a92637547..be32bba08 100644 --- a/tests/homeassistant/test_module.py +++ b/tests/homeassistant/test_module.py @@ -1,8 +1,10 @@ """Test Homeassistant module.""" +import asyncio from pathlib import Path from unittest.mock import patch +from supervisor.const import CoreState from supervisor.coresys import CoreSys from supervisor.docker.interface import DockerInterface from supervisor.homeassistant.secrets import HomeAssistantSecrets @@ -10,6 +12,7 @@ from supervisor.homeassistant.secrets import HomeAssistantSecrets async def test_load(coresys: CoreSys, tmp_supervisor_data: Path): """Test homeassistant module load.""" + client = coresys.homeassistant.websocket._client # pylint: disable=protected-access with open(tmp_supervisor_data / "homeassistant" / "secrets.yaml", "w") as secrets: secrets.write("hello: world\n") @@ -24,3 +27,11 @@ async def test_load(coresys: CoreSys, tmp_supervisor_data: Path): attach.assert_called_once() assert coresys.homeassistant.secrets.secrets == {"hello": "world"} + + coresys.core.state = CoreState.SETUP + await coresys.homeassistant.websocket.async_send_message({"lorem": "ipsum"}) + client.async_send_command.assert_not_called() + + coresys.core.state = CoreState.RUNNING + await asyncio.sleep(0) + assert client.async_send_command.call_args_list[0][0][0] == {"lorem": "ipsum"} diff --git a/tests/homeassistant/test_websocket.py b/tests/homeassistant/test_websocket.py index 823337a98..a5b636a63 100644 --- a/tests/homeassistant/test_websocket.py +++ b/tests/homeassistant/test_websocket.py @@ -1,9 +1,11 @@ """Test websocket.""" # pylint: disable=protected-access, import-error +import asyncio import logging from awesomeversion import AwesomeVersion +from supervisor.const import CoreState from supervisor.coresys import CoreSys from supervisor.homeassistant.const import WSEvent, WSType @@ -48,3 +50,36 @@ async def test_send_command_old_core_version(coresys: CoreSys, caplog): "test", {"lorem": "ipsum"} ) client.async_send_command.assert_not_called() + + +async def test_send_message_during_startup(coresys: CoreSys): + """Test websocket messages queue during startup.""" + client = coresys.homeassistant.websocket._client + await coresys.homeassistant.websocket.load() + coresys.core.state = CoreState.SETUP + + await coresys.homeassistant.websocket.async_supervisor_update_event( + "test", {"lorem": "ipsum"} + ) + client.async_send_command.assert_not_called() + + coresys.core.state = CoreState.RUNNING + await asyncio.sleep(0) + + assert client.async_send_command.call_count == 2 + assert client.async_send_command.call_args_list[0][0][0] == { + "type": WSType.SUPERVISOR_EVENT, + "data": { + "event": WSEvent.SUPERVISOR_UPDATE, + "update_key": "test", + "data": {"lorem": "ipsum"}, + }, + } + assert client.async_send_command.call_args_list[1][0][0] == { + "type": WSType.SUPERVISOR_EVENT, + "data": { + "event": WSEvent.SUPERVISOR_UPDATE, + "update_key": "info", + "data": {"state": "running"}, + }, + }