Separate startup event from update check event (#4425)

* Separate startup event from update check event

* Add a queue for messages sent during startup
This commit is contained in:
Mike Degatano 2023-07-06 12:45:37 -04:00 committed by GitHub
parent 3b38047fd4
commit 96d5fc244e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 105 additions and 6 deletions

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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