Home-Assistant
/
core
mirror of https://github.com/home-assistant/core.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Persist roborock maps to disk only on shutdown (#136889) 

* Persist roborock maps to disk only on shutdown

* Rename on_unload to on_stop

* Spawn 1 executor thread and block writes to disk

* Update tests/components/roborock/test_image.py

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>

* Use config entry setup instead of component setup

---------

Co-authored-by: Joost Lekkerkerker <joostlek@outlook.com>
pull/137390/head
Allen Porter 2025-01-31 10:45:01 -08:00 committed by Paulus Schoutsen
parent a74328e600
commit 9cd48dd452
7 changed files with 79 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
homeassistant/components/roborock/__init__.py
View File

@ -22,7 +22,7 @@ from roborock.version_a01_apis import RoborockMqttClientA01
from roborock.web_api import RoborockApiClient
from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_USERNAME
from homeassistant.const import CONF_USERNAME, EVENT_HOMEASSISTANT_STOP
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
@ -118,13 +118,21 @@ async def async_setup_entry(hass: HomeAssistant, entry: RoborockConfigEntry) ->
)
valid_coordinators = RoborockCoordinators(v1_coords, a01_coords)
async def on_unload() -> None:
release_tasks = set()
for coordinator in valid_coordinators.values():
release_tasks.add(coordinator.release())
await asyncio.gather(*release_tasks)
async def on_stop(_: Any) -> None:
_LOGGER.debug("Shutting down roborock")
await asyncio.gather(
*(
coordinator.async_shutdown()
for coordinator in valid_coordinators.values()
)
)
entry.async_on_unload(on_unload)
entry.async_on_unload(
hass.bus.async_listen_once(
EVENT_HOMEASSISTANT_STOP,
on_stop,
)
)
entry.runtime_data = valid_coordinators
await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
@ -209,7 +217,7 @@ async def setup_device_v1(
try:
await coordinator.async_config_entry_first_refresh()
except ConfigEntryNotReady as ex:
await coordinator.release()
await coordinator.async_shutdown()
if isinstance(coordinator.api, RoborockMqttClientV1):
_LOGGER.warning(
"Not setting up %s because the we failed to get data for the first time using the online client. "

18
homeassistant/components/roborock/coordinator.py
View File

@ -2,6 +2,7 @@
from __future__ import annotations
import asyncio
from datetime import timedelta
import logging
@ -116,10 +117,14 @@ class RoborockDataUpdateCoordinator(DataUpdateCoordinator[DeviceProp]):
# Right now this should never be called if the cloud api is the primary api,
# but in the future if it is, a new else should be added.
async def release(self) -> None:
"""Disconnect from API."""
await self.api.async_release()
await self.cloud_api.async_release()
async def async_shutdown(self) -> None:
"""Shutdown the coordinator."""
await super().async_shutdown()
await asyncio.gather(
self.map_storage.flush(),
self.api.async_release(),
self.cloud_api.async_release(),
)
async def _update_device_prop(self) -> None:
"""Update device properties."""
@ -226,8 +231,9 @@ class RoborockDataUpdateCoordinatorA01(
) -> dict[RoborockDyadDataProtocol | RoborockZeoProtocol, StateType]:
return await self.api.update_values(self.request_protocols)
async def release(self) -> None:
"""Disconnect from API."""
async def async_shutdown(self) -> None:
"""Shutdown the coordinator on config entry unload."""
await super().async_shutdown()
await self.api.async_release()
@cached_property

10
homeassistant/components/roborock/image.py
View File

@ -157,13 +157,9 @@ class RoborockMap(RoborockCoordinatedEntityV1, ImageEntity):
)
if self.cached_map != content:
self.cached_map = content
self.config_entry.async_create_task(
self.hass,
self.coordinator.map_storage.async_save_map(
self.map_flag,
content,
),
f"{self.unique_id} map",
await self.coordinator.map_storage.async_save_map(
self.map_flag,
content,
)
return self.cached_map

20
homeassistant/components/roborock/roborock_storage.py
View File

@ -31,6 +31,7 @@ class RoborockMapStorage:
self._path_prefix = (
_storage_path_prefix(hass, entry_id) / MAPS_PATH / device_id_slug
)
self._write_queue: dict[int, bytes] = {}
async def async_load_map(self, map_flag: int) -> bytes | None:
"""Load maps from disk."""
@ -48,9 +49,22 @@ class RoborockMapStorage:
return None
async def async_save_map(self, map_flag: int, content: bytes) -> None:
"""Write map if it should be updated."""
filename = self._path_prefix / f"{map_flag}{MAP_FILENAME_SUFFIX}"
await self._hass.async_add_executor_job(self._save_map, filename, content)
"""Save the map to a pending write queue."""
self._write_queue[map_flag] = content
async def flush(self) -> None:
"""Flush all maps to disk."""
_LOGGER.debug("Flushing %s maps to disk", len(self._write_queue))
queue = self._write_queue.copy()
def _flush_all() -> None:
for map_flag, content in queue.items():
filename = self._path_prefix / f"{map_flag}{MAP_FILENAME_SUFFIX}"
self._save_map(filename, content)
await self._hass.async_add_executor_job(_flush_all)
self._write_queue.clear()
def _save_map(self, filename: Path, content: bytes) -> None:
"""Write the map to disk."""

10
tests/components/roborock/conftest.py
View File

@ -19,9 +19,9 @@ from homeassistant.components.roborock.const import (
CONF_USER_DATA,
DOMAIN,
)
from homeassistant.config_entries import ConfigEntryState
from homeassistant.const import CONF_USERNAME, Platform
from homeassistant.core import HomeAssistant
from homeassistant.setup import async_setup_component
from .mock_data import (
BASE_URL,
@ -207,13 +207,13 @@ async def setup_entry(
) -> Generator[MockConfigEntry]:
"""Set up the Roborock platform."""
with patch("homeassistant.components.roborock.PLATFORMS", platforms):
assert await async_setup_component(hass, DOMAIN, {})
await hass.config_entries.async_setup(mock_roborock_entry.entry_id)
await hass.async_block_till_done()
yield mock_roborock_entry
@pytest.fixture
def cleanup_map_storage(
async def cleanup_map_storage(
hass: HomeAssistant, mock_roborock_entry: MockConfigEntry
) -> Generator[pathlib.Path]:
"""Test cleanup, remove any map storage persisted during the test."""
@ -225,4 +225,8 @@ def cleanup_map_storage(
pathlib.Path(hass.config.path(tmp_path)) / mock_roborock_entry.entry_id
)
yield storage_path
# We need to first unload the config entry because unloading it will
# persist any unsaved maps to storage.
if mock_roborock_entry.state is ConfigEntryState.LOADED:
await hass.config_entries.async_unload(mock_roborock_entry.entry_id)
shutil.rmtree(str(storage_path), ignore_errors=True)

28
tests/components/roborock/test_image.py
View File

@ -12,6 +12,7 @@ from roborock import RoborockException
from vacuum_map_parser_base.map_data import ImageConfig, ImageData
from homeassistant.components.roborock import DOMAIN
from homeassistant.config_entries import ConfigEntryState
from homeassistant.const import Platform
from homeassistant.core import HomeAssistant
from homeassistant.setup import async_setup_component
@ -120,7 +121,7 @@ async def test_load_stored_image(
MAP_DATA.image.data.save(img_byte_arr, format="PNG")
img_bytes = img_byte_arr.getvalue()
# Load the image on demand, which should ensure it is cached on disk
# Load the image on demand, which should queue it to be cached on disk
client = await hass_client()
resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
assert resp.status == HTTPStatus.OK
@ -151,22 +152,25 @@ async def test_fail_to_save_image(
caplog: pytest.LogCaptureFixture,
) -> None:
"""Test that we gracefully handle a oserror on saving an image."""
# Reload the config entry so that the map is saved in storage and entities exist.
await async_setup_component(hass, DOMAIN, {})
await hass.async_block_till_done()
# Ensure that map is still working properly.
assert hass.states.get("image.roborock_s7_maxv_upstairs") is not None
client = await hass_client()
resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
# Test that we can get the image and it correctly serialized and unserialized.
assert resp.status == HTTPStatus.OK
with patch(
"homeassistant.components.roborock.roborock_storage.Path.write_bytes",
side_effect=OSError,
):
await async_setup_component(hass, DOMAIN, {})
await hass.async_block_till_done()
await hass.config_entries.async_unload(mock_roborock_entry.entry_id)
assert "Unable to write map file" in caplog.text
# Ensure that map is still working properly.
assert hass.states.get("image.roborock_s7_maxv_upstairs") is not None
client = await hass_client()
resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
# Test that we can get the image and it correctly serialized and unserialized.
assert resp.status == HTTPStatus.OK
assert "Unable to write map file" in caplog.text
# Config entry is unloaded successfully
assert mock_roborock_entry.state is ConfigEntryState.NOT_LOADED
async def test_fail_to_load_image(

8
tests/components/roborock/test_init.py
View File

@ -183,6 +183,10 @@ async def test_remove_from_hass(
resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
assert resp.status == HTTPStatus.OK
assert not cleanup_map_storage.exists()
# Flush to disk
await hass.config_entries.async_unload(setup_entry.entry_id)
assert cleanup_map_storage.exists()
paths = list(cleanup_map_storage.walk())
assert len(paths) == 3 # One map image and two directories
@ -209,6 +213,10 @@ async def test_oserror_remove_image(
resp = await client.get("/api/image_proxy/image.roborock_s7_maxv_upstairs")
assert resp.status == HTTPStatus.OK
# Image content is saved when unloading
assert not cleanup_map_storage.exists()
await hass.config_entries.async_unload(setup_entry.entry_id)
assert cleanup_map_storage.exists()
paths = list(cleanup_map_storage.walk())
assert len(paths) == 3 # One map image and two directories