Separate startup event from update check event (#4425)

* Separate startup event from update check event

* Add a queue for messages sent during startup

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,
+]

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

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()),
             ]

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
 

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"}

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"},
+        },
+    }